From d4bbf728ca52ce23a549a299e229666e734ffb7c Mon Sep 17 00:00:00 2001 From: Junkai-Wu Date: Sat, 14 Feb 2026 12:27:58 +0800 Subject: [PATCH] v4.4 tag release update. (#3032) --- CHANGELOG.md | 30 +- README.md | 30 +- .../75_blackwell_grouped_gemm.cu | 16 +- .../75_blackwell_grouped_gemm_block_scaled.cu | 16 +- .../75_blackwell_grouped_gemm/CMakeLists.txt | 8 +- .../CuTeDSL/ampere/call_bypass_dlpack.py | 3 +- .../python/CuTeDSL/ampere/call_from_jit.py | 6 +- .../CuTeDSL/ampere/cooperative_launch.py | 627 +++ .../python/CuTeDSL/ampere/elementwise_add.py | 15 +- .../ampere/elementwise_add_autotune.py | 372 ++ .../CuTeDSL/ampere/elementwise_apply.py | 7 +- .../CuTeDSL/ampere/flash_attention_v2.py | 12 +- .../python/CuTeDSL/ampere/hstu_attention.py | 22 +- examples/python/CuTeDSL/ampere/inline_ptx.py | 15 +- examples/python/CuTeDSL/ampere/sgemm.py | 6 +- .../python/CuTeDSL/ampere/smem_allocator.py | 3 +- .../python/CuTeDSL/ampere/tensorop_gemm.py | 6 +- .../blockwise_gemm/blockwise_gemm.py | 46 +- .../blockwise_gemm/contiguous_grouped_gemm.py | 47 +- .../blockwise_gemm/masked_grouped_gemm.py | 47 +- .../dense_blockscaled_gemm_persistent.py | 1212 +++-- .../dense_blockscaled_gemm_persistent_amax.py | 2576 ++++++++++ .../python/CuTeDSL/blackwell/dense_gemm.py | 21 +- .../dense_gemm_alpha_beta_persistent.py | 20 +- .../blackwell/dense_gemm_persistent.py | 870 ++-- .../dense_gemm_persistent_dynamic.py | 11 +- .../blackwell/dense_gemm_software_pipeline.py | 28 +- .../activation_custom_epilogue_dense_gemm.py | 754 +++ .../epilogue/common_dense_gemm_efc.py | 2011 ++++++++ .../CuTeDSL/blackwell/epilogue/common_efc.py | 1512 ++++++ .../epilogue/custom_epilogue_dense_gemm.py | 625 +++ .../synthetic_custom_epilogue_dense_gemm.py | 403 ++ examples/python/CuTeDSL/blackwell/fmha.py | 58 +- examples/python/CuTeDSL/blackwell/fmha_bwd.py | 16 +- .../blackwell/grouped_blockscaled_gemm.py | 630 +-- .../python/CuTeDSL/blackwell/grouped_gemm.py | 688 +-- .../blackwell/mamba2_ssd/mamba2_ssd.py | 181 +- .../mixed_input_fmha_decode.py | 1435 +++--- .../mixed_input_fmha_prefill_d256.py | 2030 ++++++++ .../mixed_input_fmha_prefill_d512.py | 2172 ++++++++ .../mixed_input_fmha/prefill_helpers.py | 400 ++ .../grouped_mixed_input_gemm.py | 2526 ++++++++++ .../grouped_mixed_input_gemm_acc_scale.py | 2502 ++++++++++ .../mixed_input_gemm/mixed_input_gemm.py | 2321 +++++++++ .../mixed_input_host_utils.py | 506 ++ .../CuTeDSL/blackwell/mla/mla_decode_fp16.py | 4373 +++++++++++++++++ .../CuTeDSL/blackwell/mla/mla_decode_fp8.py | 4341 ++++++++++++++++ .../CuTeDSL/blackwell/mla/mla_helpers.py | 304 ++ ...sm103_dense_blockscaled_gemm_persistent.py | 78 +- .../blackwell/tutorial_gemm/fp16_gemm_0.py | 45 +- .../blackwell/tutorial_gemm/fp16_gemm_1.py | 34 +- .../blackwell/tutorial_gemm/nvfp4_gemm_0.py | 778 +++ .../blackwell/tutorial_gemm/nvfp4_gemm_1.py | 934 ++++ .../CuTeDSL/blackwell/tutorial_gemm/utils.py | 366 ++ .../CuTeDSL/blackwell_geforce/dense_gemm.py | 20 +- .../python/CuTeDSL/cute/ffi/jit_argument.py | 3 +- .../python/CuTeDSL/cute/torch_fake_tensor.py | 3 +- .../cute/tvm_ffi/aot_use_in_cpp_bundle.cpp | 2 +- .../cute/tvm_ffi/aot_use_in_cpp_bundle.sh | 2 +- .../CuTeDSL/distributed/all_reduce_tma.py | 6 +- .../blackwell/dense_block_scaled_gemm.py | 28 +- .../experimental/blackwell/dense_gemm.py | 35 +- .../experimental/blackwell/dense_gemm_2sm.py | 13 +- .../blackwell/dense_gemm_cute_pipeline.py | 587 ++- .../blackwell/dense_gemm_ptr_array.py | 11 +- .../python/CuTeDSL/helpers/sparse_utils.py | 457 -- .../CuTeDSL/helpers/test_sparse_utils.py | 104 - examples/python/CuTeDSL/hopper/dense_gemm.py | 13 +- .../CuTeDSL/hopper/dense_gemm_persistent.py | 10 +- examples/python/CuTeDSL/hopper/fmha.py | 25 +- .../CuTeDSL/jax/cutlass_call_sharding.py | 69 +- .../CuTeDSL/notebooks/hello_world.ipynb | 4 +- .../cutlass/arch/grid_dependency_control.h | 2 + ...m100_implicit_gemm_tma_warpspecialized.hpp | 6 +- ...100_mma_array_warpspecialized_emulated.hpp | 5 + ...ecialized_interleaved_complex_emulated.hpp | 5 + ...rpspecialized_interleaved_complex_tf32.hpp | 5 + .../sm100_mma_warpspecialized_emulated.hpp | 5 + ...ecialized_interleaved_complex_emulated.hpp | 5 + ...rpspecialized_interleaved_complex_tf32.hpp | 5 + ...ay_tma_warpspecialized_input_transform.hpp | 37 +- ...mm_tma_warpspecialized_input_transform.hpp | 38 +- .../kernel/sm100_tile_scheduler_group.hpp | 3 - ..._array_tma_warpspecialized_cooperative.hpp | 3 - ...emm_array_tma_warpspecialized_pingpong.hpp | 3 - .../gemm/kernel/sm90_tile_scheduler_group.hpp | 42 +- .../kernel/sm90_tile_scheduler_stream_k.hpp | 5 +- .../gemm/kernel/tile_scheduler_params.h | 46 +- include/cutlass/platform/platform.h | 12 +- media/docs/pythonDSL/cute_dsl.rst | 1 + .../pythonDSL/cute_dsl_general/debugging.rst | 3 - .../dsl_ahead_of_time_compilation.rst | 2 +- .../cute_dsl_general/dsl_code_generation.rst | 242 +- .../cute_dsl_general/dsl_compilation.png | Bin 0 -> 288213 bytes .../cute_dsl_general/dsl_introduction.rst | 24 +- .../pythonDSL/cute_dsl_general/resources.rst | 28 + media/docs/pythonDSL/overview.rst | 2 +- .../CuTeDSL/cutlass/base_dsl/ast_helpers.py | 12 +- .../cutlass/base_dsl/ast_preprocessor.py | 32 +- python/CuTeDSL/cutlass/base_dsl/dsl.py | 11 +- .../base_dsl/tvm_ffi_builder/mlir_builder.py | 57 - python/CuTeDSL/cutlass/cute/algorithm.py | 88 +- .../cutlass/cute/arch/nvvm_wrappers.py | 41 +- python/CuTeDSL/cutlass/cute/atom.py | 96 +- .../cutlass/cute/experimental/README.md | 3 + .../cutlass/cute/experimental/utils.py | 85 +- .../cutlass/cute/nvgpu/tcgen05/__init__.py | 1 + .../cutlass/cute/nvgpu/tcgen05/copy.py | 107 + .../CuTeDSL/cutlass/cute/nvgpu/warp/copy.py | 15 +- python/CuTeDSL/cutlass/cute/runtime.py | 12 +- .../cutlass_dsl/cutlass_ast_decorators.py | 14 +- .../cutlass/cutlass_dsl/tvm_ffi_provider.py | 17 +- .../dynamic_persistent_tile_scheduler.py | 4 +- python/CuTeDSL/cutlass/utils/gemm/sm100.py | 53 +- python/CuTeDSL/cutlass/utils/hardware_info.py | 25 +- python/CuTeDSL/pyproject.toml | 2 +- python/CuTeDSL/requirements-cu13.txt | 3 + python/CuTeDSL/requirements.txt | 2 +- python/CuTeDSL/setup.sh | 85 + python/cutlass_library/heuristics_provider.py | 109 +- .../device/sm100_tensorop_gemm/CMakeLists.txt | 1 + .../extra_tests/CMakeLists.txt | 95 + .../sm100_gemm_bf16_bf16_f32_tensor_op_f32.cu | 310 ++ ...m100_gemm_f16_f16_f16_tensor_op_f32_b2b.cu | 257 + ...gemm_f16_f16_f16_tensor_op_f32_stream_k.cu | 235 + ...f16_f16_tensor_op_f32_swap_ab_bias_relu.cu | 168 + .../sm100_gemm_f16_f16_f32_tensor_op_f32.cu | 96 + .../sm100_gemm_f16_f16_f32_void_f16_stage.cu | 233 + .../sm100_gemm_f32_f32_f32_void_f32_stage.cu | 234 + ...4_f4_f32_tensor_op_f32_runtime_datatype.cu | 154 + ...6_f6_f32_tensor_op_f32_runtime_datatype.cu | 154 + ...8_f4_f32_tensor_op_f32_runtime_datatype.cu | 107 + .../sm100_gemm_f8_f8_f32_void_f8_stage.cu | 234 + ...0_gemm_f8_f8_f8_tensor_op_f32_bias_gelu.cu | 319 ++ ..._f8_f8_tensor_op_f32_bias_gelu_amax_aux.cu | 328 ++ ...0_gemm_f8_f8_f8_tensor_op_f32_bias_relu.cu | 365 ++ ...f8_f8_f8_tensor_op_f32_runtime_datatype.cu | 295 ++ ...f8_f8_f8_tensor_op_s32_batch_alpha_beta.cu | 219 + ...0_gemm_i8_i8_i8_tensor_op_s32_bias_relu.cu | 278 ++ ...8_i8_i8_tensor_op_s32_vector_alpha_beta.cu | 343 ++ 140 files changed, 41624 insertions(+), 3691 deletions(-) create mode 100644 examples/python/CuTeDSL/ampere/cooperative_launch.py create mode 100644 examples/python/CuTeDSL/ampere/elementwise_add_autotune.py create mode 100644 examples/python/CuTeDSL/blackwell/dense_blockscaled_gemm_persistent_amax.py create mode 100644 examples/python/CuTeDSL/blackwell/epilogue/activation_custom_epilogue_dense_gemm.py create mode 100644 examples/python/CuTeDSL/blackwell/epilogue/common_dense_gemm_efc.py create mode 100644 examples/python/CuTeDSL/blackwell/epilogue/common_efc.py create mode 100644 examples/python/CuTeDSL/blackwell/epilogue/custom_epilogue_dense_gemm.py create mode 100644 examples/python/CuTeDSL/blackwell/epilogue/synthetic_custom_epilogue_dense_gemm.py create mode 100644 examples/python/CuTeDSL/blackwell/mixed_input_fmha/mixed_input_fmha_prefill_d256.py create mode 100644 examples/python/CuTeDSL/blackwell/mixed_input_fmha/mixed_input_fmha_prefill_d512.py create mode 100644 examples/python/CuTeDSL/blackwell/mixed_input_fmha/prefill_helpers.py create mode 100644 examples/python/CuTeDSL/blackwell/mixed_input_gemm/grouped_mixed_input_gemm.py create mode 100644 examples/python/CuTeDSL/blackwell/mixed_input_gemm/grouped_mixed_input_gemm_acc_scale.py create mode 100644 examples/python/CuTeDSL/blackwell/mixed_input_gemm/mixed_input_gemm.py create mode 100644 examples/python/CuTeDSL/blackwell/mixed_input_gemm/mixed_input_host_utils.py create mode 100644 examples/python/CuTeDSL/blackwell/mla/mla_decode_fp16.py create mode 100644 examples/python/CuTeDSL/blackwell/mla/mla_decode_fp8.py create mode 100644 examples/python/CuTeDSL/blackwell/mla/mla_helpers.py create mode 100644 examples/python/CuTeDSL/blackwell/tutorial_gemm/nvfp4_gemm_0.py create mode 100644 examples/python/CuTeDSL/blackwell/tutorial_gemm/nvfp4_gemm_1.py create mode 100644 examples/python/CuTeDSL/blackwell/tutorial_gemm/utils.py delete mode 100644 examples/python/CuTeDSL/helpers/sparse_utils.py delete mode 100644 examples/python/CuTeDSL/helpers/test_sparse_utils.py create mode 100644 media/docs/pythonDSL/cute_dsl_general/dsl_compilation.png create mode 100644 media/docs/pythonDSL/cute_dsl_general/resources.rst create mode 100644 python/CuTeDSL/requirements-cu13.txt create mode 100755 python/CuTeDSL/setup.sh create mode 100644 test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/CMakeLists.txt create mode 100644 test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_bf16_bf16_f32_tensor_op_f32.cu create mode 100644 test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_f16_f16_f16_tensor_op_f32_b2b.cu create mode 100644 test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_f16_f16_f16_tensor_op_f32_stream_k.cu create mode 100644 test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_f16_f16_f16_tensor_op_f32_swap_ab_bias_relu.cu create mode 100644 test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_f16_f16_f32_tensor_op_f32.cu create mode 100644 test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_f16_f16_f32_void_f16_stage.cu create mode 100644 test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_f32_f32_f32_void_f32_stage.cu create mode 100644 test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_f4_f4_f32_tensor_op_f32_runtime_datatype.cu create mode 100644 test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_f6_f6_f32_tensor_op_f32_runtime_datatype.cu create mode 100644 test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_f8_f4_f32_tensor_op_f32_runtime_datatype.cu create mode 100644 test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_f8_f8_f32_void_f8_stage.cu create mode 100644 test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_f8_f8_f8_tensor_op_f32_bias_gelu.cu create mode 100644 test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_f8_f8_f8_tensor_op_f32_bias_gelu_amax_aux.cu create mode 100644 test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_f8_f8_f8_tensor_op_f32_bias_relu.cu create mode 100644 test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_f8_f8_f8_tensor_op_f32_runtime_datatype.cu create mode 100644 test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_f8_f8_f8_tensor_op_s32_batch_alpha_beta.cu create mode 100644 test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_i8_i8_i8_tensor_op_s32_bias_relu.cu create mode 100644 test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_i8_i8_i8_tensor_op_s32_vector_alpha_beta.cu diff --git a/CHANGELOG.md b/CHANGELOG.md index af35b23a..17048ffa 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,7 +2,7 @@ # CUTLASS 4.x -## [4.4.0](https://github.com/NVIDIA/cutlass/tree/main) (2026-01-23) +## [4.4.0](https://github.com/NVIDIA/cutlass/releases/tag/v4.4.0) (2026-02-14) ### CuTe DSL * New features @@ -18,6 +18,7 @@ + New pipeline abstraction with convenience wrappers + New Partition ops to simplify partitioning logic. + Device-side TMA descriptor allocation, initialization, and management + + These examples can be found here https://github.com/NVIDIA/cutlass/tree/main/examples/python/CuTeDSL/experimental - Ahead of Time (AoT) compilation is now available! + Refer to files under https://github.com/NVIDIA/cutlass/tree/main/examples/python/CuTeDSL/cute/export for example usage - JAX support - you can now use CuTeDSL along with JAX @@ -26,21 +27,32 @@ + cutlass.__version__ for a string representation of DSL version + cutlass.CUDA_VERSION for a version class to tell the CUDA version used for DSL - Added CopyDsmemStoreOp to store data to distributed shared memory with explicit synchronization. + - Grouped GEMM example now supports device-only problem shapes. + - We allow grid carve-out without problem shapes being available on host. + - Tma+LdMatrix features for loading+unpacking narrow-width types (refer to mixed_input_fmha_decode.py for example usage). + - It is possible now to have customized epilogue fusion for persistent dense GEMM through a Python Epilogue Fusion Configuration (EFC) function, somewhat similar to CUTLASS C++ EVT. It also provides a PyTorch evaluator to compare the results. * More examples of authorizing peak-performance kernels - [SM103 batched 3xFP4 blockscaled GEMM kernel](https://github.com/NVIDIA/cutlass/tree/main/examples/python/CuTeDSL/blackwell/sm103_dense_blockscaled_gemm_persistent.py) + - Mixed input FMHA decode example with support for int4 KV (int8 KV supported in 4.3) + - New acc_scale grouped mixed input gemm kernel variant is introduced to deliver better performance for decoding cases. + - All mixed_input_gemm examples are moved into a separate folder `mixed_input_gemm`. Common utility functions are also extracted into mixed_input_host_utils.py under the same folder. * Bug fixing and improvements - Fixed an issue that both branches of if are executed - Fixed `cute.printf` with f-string - - Fixed an issue that cutlass.cuda.initialize_cuda_context() silently kills python + - Fixed an indexing issue of scalar tensor + - Fixed small K reference check error for cta_tile_n = 256 case with overlapping accumulator optimization in [Blackwell SM100 persistent dense blockscaled GEMM with static scheduling](https://github.com/NVIDIA/cutlass/tree/main/examples/python/CuTeDSL/blackwell/dense_blockscaled_gemm_persistent.py). * API changes - Deprecate get_num_tmem_alloc_cols from blackwell_helpers.py. Use the one from tmem_allocator.py instead. - Deprecate SM100_TMEM_CAPACITY_COLUMNS and SM100_TMEM_MIN_ALLOC_COLUMNS. - LdMatrix16x16x8bOp and StMatrix16x8x8bOp now require explicit transpose=True when calling __init__, to avoid ambiguity in data transposition. - LdMatrix16x16x8bOp copy traits updated to be faithful to PTX without permutations. Permuted variant is renamed to LdMatrix16x8x8bOp. + - Grouped GEMM example takes the argument --host_problem_shape_available. If the argument is provided, grid is carved out based upon the host problem shapes, otherwise, we launch maximum possible SMs. + - hardware_info.get_max_active_cluster support pass in specific stream to query. Useful for green context based SM partition. - group_bulk_copy_modes in async bulk copy example is now deprecated, use group_modes directly instead. + - Deprecate nvvm wrapper from using nvvm enum, use str instead. - cute.arch.calc_packed_f32x2_op default enable ftz to default disable ftz - In CuTe DSL with CTK 13.1, following APIs in cutlass.cute.arch now require string literal instead of enum as argument: + fence_proxy @@ -58,6 +70,12 @@ + store + load +* Advanced compiler control +Use 'Advanced compiler control' for mixed input gemm examples for better performance. +Advanced compiler control is an experimental feature of CUDA compiler. The controls file contains internal compiler settings tuned for specific kernels with a specific version of CUDA toolkit to get better GPU kernel code. More details and documentation on how to create these controls files will be provided in future CUDA toolkit release. + +Note: The advanced compiler control file is not expected to work for kernels that it was not tuned for. There is no compatibility guarantee, and the controls file will not work for CUDA toolkit with a different version. + ### CUTLASS C++ * Add Hopper e2m1 to fp32 optimized conversion and e2m1 * TF32 tensor core GEMM. - Set MmaType to tfloat32_t for FP32 mode. @@ -86,12 +104,18 @@ - Fix memory fence for clc scheduler in Blackwell SM120 pingpong kernel. - Fix missing SMEM alignment in Blackwell SM120 scale factors. - Fix a PDL issue for grouped gemm. + - Fix divide-by-zero issue in canimplement for sm100 implicit gemm kernels. + - Fix cluster swizzle for Grouped GEMMs. + + Move host-side swizzling heuristics to device. + + Apply swizzle per group based on problem shape and max swizzle size. + + Improve examples and unit tests. * Fix some profiler issues: - Refactor L1 functional test generation logic to reduce the L1 test cases to avoid timeout. - Fix a core dump issue for nvfp4 grouped GEMM kernel. - Fix inconsistent GEMM verification logic. - Rework grouped gemm verification logic for different types. -* Fix some broken links under `media/docs`. + - Fix api break change in libheuristics. +* Fix some failed links under `media/docs`. * Various improvements and fixes from the community and CUTLASS team. Thanks to everyone who submitted PRs! * Optimal code generation with CUDA toolkit versions 13.1. diff --git a/README.md b/README.md index 00d3ffac..869c885d 100644 --- a/README.md +++ b/README.md @@ -3,7 +3,7 @@ # CUTLASS 4.4.0 -_CUTLASS 4.4.0 - Jan 2026_ +_CUTLASS 4.4.0 - Feb 2026_ CUTLASS is a collection of abstractions for implementing high-performance matrix-matrix multiplication (GEMM) and related computations at all levels and scales within CUDA. It incorporates strategies for @@ -59,6 +59,7 @@ To get started quickly - please refer : + New pipeline abstraction with convenience wrappers + New Partition ops to simplify partitioning logic. + Device-side TMA descriptor allocation, initialization, and management + + These examples can be found here https://github.com/NVIDIA/cutlass/tree/main/examples/python/CuTeDSL/experimental - Ahead of Time (AoT) compilation is now available! + Refer to files under https://github.com/NVIDIA/cutlass/tree/main/examples/python/CuTeDSL/cute/export for example usage - JAX support - you can now use CuTeDSL along with JAX @@ -67,21 +68,32 @@ To get started quickly - please refer : + cutlass.__version__ for a string representation of DSL version + cutlass.CUDA_VERSION for a version class to tell the CUDA version used for DSL - Added CopyDsmemStoreOp to store data to distributed shared memory with explicit synchronization. + - Grouped GEMM example now supports device-only problem shapes. + - We allow grid carve-out without problem shapes being available on host. + - Tma+LdMatrix features for loading+unpacking narrow-width types (refer to mixed_input_fmha_decode.py for example usage). + - It is possible now to have customized epilogue fusion for persistent dense GEMM through a Python Epilogue Fusion Configuration (EFC) function, somewhat similar to CUTLASS C++ EVT. It also provides a PyTorch evaluator to compare the results. * More examples of authorizing peak-performance kernels - [SM103 batched 3xFP4 blockscaled GEMM kernel](https://github.com/NVIDIA/cutlass/tree/main/examples/python/CuTeDSL/blackwell/sm103_dense_blockscaled_gemm_persistent.py) + - Mixed input FMHA decode example with support for int4 KV (int8 KV supported in 4.3) + - New acc_scale grouped mixed input gemm kernel variant is introduced to deliver better performance for decoding cases. + - All mixed_input_gemm examples are moved into a separate folder `mixed_input_gemm`. Common utility functions are also extracted into mixed_input_host_utils.py under the same folder. * Bug fixing and improvements - Fixed an issue that both branches of if are executed - Fixed `cute.printf` with f-string - - Fixed an issue that cutlass.cuda.initialize_cuda_context() silently kills python + - Fixed an indexing issue of scalar tensor + - Fixed small K reference check error for cta_tile_n = 256 case with overlapping accumulator optimization in [Blackwell SM100 persistent dense blockscaled GEMM with static scheduling](https://github.com/NVIDIA/cutlass/tree/main/examples/python/CuTeDSL/blackwell/dense_blockscaled_gemm_persistent.py). * API changes - Deprecate get_num_tmem_alloc_cols from blackwell_helpers.py. Use the one from tmem_allocator.py instead. - Deprecate SM100_TMEM_CAPACITY_COLUMNS and SM100_TMEM_MIN_ALLOC_COLUMNS. - LdMatrix16x16x8bOp and StMatrix16x8x8bOp now require explicit transpose=True when calling __init__, to avoid ambiguity in data transposition. - LdMatrix16x16x8bOp copy traits updated to be faithful to PTX without permutations. Permuted variant is renamed to LdMatrix16x8x8bOp. + - Grouped GEMM example takes the argument --host_problem_shape_available. If the argument is provided, grid is carved out based upon the host problem shapes, otherwise, we launch maximum possible SMs. + - hardware_info.get_max_active_cluster support pass in specific stream to query. Useful for green context based SM partition. - group_bulk_copy_modes in async bulk copy example is now deprecated, use group_modes directly instead. + - Deprecate nvvm wrapper from using nvvm enum, use str instead. - cute.arch.calc_packed_f32x2_op default enable ftz to default disable ftz - In CuTe DSL with CTK 13.1, following APIs in cutlass.cute.arch now require string literal instead of enum as argument: + fence_proxy @@ -99,6 +111,12 @@ To get started quickly - please refer : + store + load +* Advanced compiler control +Use 'Advanced compiler control' for mixed input gemm examples for better performance. +Advanced compiler control is an experimental feature of CUDA compiler. The controls file contains internal compiler settings tuned for specific kernels with a specific version of CUDA toolkit to get better GPU kernel code. More details and documentation on how to create these controls files will be provided in future CUDA toolkit release. + +Note: The advanced compiler control file is not expected to work for kernels that it was not tuned for. There is no compatibility guarantee, and the controls file will not work for CUDA toolkit with a different version. + ## CUTLASS C++ * Add Hopper e2m1 to fp32 optimized conversion and e2m1 * TF32 tensor core GEMM. - Set MmaType to tfloat32_t for FP32 mode. @@ -127,12 +145,18 @@ To get started quickly - please refer : - Fix memory fence for clc scheduler in Blackwell SM120 pingpong kernel. - Fix missing SMEM alignment in Blackwell SM120 scale factors. - Fix a PDL issue for grouped gemm. + - Fix divide-by-zero issue in canimplement for sm100 implicit gemm kernels. + - Fix cluster swizzle for Grouped GEMMs. + + Move host-side swizzling heuristics to device. + + Apply swizzle per group based on problem shape and max swizzle size. + + Improve examples and unit tests. * Fix some profiler issues: - Refactor L1 functional test generation logic to reduce the L1 test cases to avoid timeout. - Fix a core dump issue for nvfp4 grouped GEMM kernel. - Fix inconsistent GEMM verification logic. - Rework grouped gemm verification logic for different types. -* Fix some broken links under `media/docs`. + - Fix api break change in libheuristics. +* Fix some failed links under `media/docs`. Note: CUTLASS 4.x builds are known to be down on Windows platforms for all CUDA toolkits. CUTLASS team is working on a fix. diff --git a/examples/75_blackwell_grouped_gemm/75_blackwell_grouped_gemm.cu b/examples/75_blackwell_grouped_gemm/75_blackwell_grouped_gemm.cu index cde24323..52c06507 100644 --- a/examples/75_blackwell_grouped_gemm/75_blackwell_grouped_gemm.cu +++ b/examples/75_blackwell_grouped_gemm/75_blackwell_grouped_gemm.cu @@ -251,6 +251,8 @@ struct Options { dim3 cluster_shape = dim3(4,2,1); dim3 cluster_shape_fallback = dim3(2,1,1); RasterOrderOptions raster_order = RasterOrderOptions::AlongM; + char raster_char = 'M'; + int swizzle = 1; int max_sm_count = INT_MAX; std::string benchmark_path; std::vector problem_sizes_host; @@ -294,7 +296,6 @@ struct Options { randomize_problems(cmd); } - char raster_char; cmd.get_cmd_line_argument("raster", raster_char); if (raster_char == 'N' || raster_char == 'n') { @@ -303,6 +304,7 @@ struct Options { else if (raster_char == 'M' || raster_char == 'm') { raster_order = RasterOrderOptions::AlongM; } + cmd.get_cmd_line_argument("swizzle", swizzle, 1); } void randomize_problems(cutlass::CommandLine &cmd) { @@ -378,7 +380,8 @@ struct Options { << " --beta= Epilogue scalar beta\n\n" << " --cluster_m= and --cluster_n= Sets the X,Y dims of the preferred cluster shape\n" << " --cluster_fallback_m= and --cluster_fallback_n= Sets the X,Y dims of the fallback cluster shape\n\n" - << " --raster= CTA Rasterization direction (N for along N, M for along M)\n\n" + << " --raster= Cluster rasterization direction (N for along N, M for along M)\n" + << " --swizzle= Cluster swizzle (swizzle up to 8 and with the nearest multiple of 2)\n\n" << " --iterations= Number of profiling iterations to perform\n\n" << " --benchmark= Executes a benchmark problem size\n" << " --max_sm_count= Run kernels using only these number of SMs\n" @@ -615,6 +618,7 @@ typename Gemm::Arguments args_from_options(Options &options, bool host_problem_s typename Gemm::GemmKernel::TileSchedulerArguments scheduler; scheduler.raster_order = options.raster_order; + scheduler.max_swizzle_size = options.swizzle; if (host_problem_shapes_available) { arguments = typename Gemm::Arguments { @@ -685,7 +689,13 @@ int run(Options &options, bool host_problem_shapes_available = true) std::cout << " " << options.problem_sizes_host.at(i); std::cout << ", " << alpha_host.at(i) << ", " << beta_host.at(i) << std::endl; } - std::cout << " Groups : " << options.groups << std::endl; + std::cout << " Groups : " << options.groups << std::endl; + + std::cout << " Cluster Shape : " << options.cluster_shape.x << "x" << options.cluster_shape.y << std::endl; + std::cout << " Cluster Fallback Shape : " << options.cluster_shape_fallback.x << "x" << options.cluster_shape_fallback.y << std::endl; + + std::cout << " Raster Order : Along-" << options.raster_char << std::endl; + std::cout << " Max Swizzle Size : " << options.swizzle << std::endl; // Instantiate CUTLASS kernel depending on templates Gemm gemm; diff --git a/examples/75_blackwell_grouped_gemm/75_blackwell_grouped_gemm_block_scaled.cu b/examples/75_blackwell_grouped_gemm/75_blackwell_grouped_gemm_block_scaled.cu index 13631595..1ff8f215 100644 --- a/examples/75_blackwell_grouped_gemm/75_blackwell_grouped_gemm_block_scaled.cu +++ b/examples/75_blackwell_grouped_gemm/75_blackwell_grouped_gemm_block_scaled.cu @@ -309,6 +309,8 @@ struct Options { dim3 cluster_shape = dim3(2,1,1); dim3 cluster_shape_fallback = dim3(2,1,1); RasterOrderOptions raster_order = RasterOrderOptions::AlongN; + char raster_char = 'N'; + int swizzle = 1; int max_sm_count = INT_MAX; std::string benchmark_path; std::vector problem_sizes_host; @@ -356,7 +358,6 @@ struct Options { randomize_problems(cmd); } - char raster_char; cmd.get_cmd_line_argument("raster", raster_char); if (raster_char == 'N' || raster_char == 'n') { @@ -365,6 +366,7 @@ struct Options { else if (raster_char == 'M' || raster_char == 'm') { raster_order = RasterOrderOptions::AlongM; } + cmd.get_cmd_line_argument("swizzle", swizzle, 1); } void randomize_problems(cutlass::CommandLine &cmd) { @@ -441,7 +443,8 @@ struct Options { << " --norm_constant= Epilogue scalar normalization constant for the output matrix\n\n" << " --cluster_m= and --cluster_n= Sets the X,Y dims of the preferred cluster shape\n" << " --cluster_fallback_m= and --cluster_fallback_n= Sets the X,Y dims of the fallback cluster shape\n\n" - << " --raster= CTA Rasterization direction (N for along N, M for along M)\n\n" + << " --raster= Cluster rasterization direction (N for along N, M for along M)\n" + << " --swizzle= Cluster swizzle (swizzle up to 8 and with the nearest multiple of 2)\n\n" << " --iterations= Number of profiling iterations to perform\n\n" << " --benchmark= Executes a benchmark problem size\n" << " --max_sm_count= Run kernels using only these number of SMs\n" @@ -722,6 +725,7 @@ typename Gemm::Arguments args_from_options(Options &options, bool host_problem_s typename Gemm::GemmKernel::TileSchedulerArguments scheduler; scheduler.raster_order = options.raster_order; + scheduler.max_swizzle_size = options.swizzle; if (host_problem_shapes_available) { arguments = typename Gemm::Arguments { @@ -813,7 +817,13 @@ int run(Options &options, bool host_problem_shapes_available = true) std::cout << " " << options.problem_sizes_host.at(i); std::cout << ", " << alpha_host.at(i) << ", " << beta_host.at(i) << std::endl; } - std::cout << " Groups : " << options.groups << std::endl; + std::cout << " Groups : " << options.groups << std::endl; + + std::cout << " Cluster Shape : " << options.cluster_shape.x << "x" << options.cluster_shape.y << std::endl; + std::cout << " Cluster Fallback Shape : " << options.cluster_shape_fallback.x << "x" << options.cluster_shape_fallback.y << std::endl; + + std::cout << " Raster Order : Along-" << options.raster_char << std::endl; + std::cout << " Max Swizzle Size : " << options.swizzle << std::endl; // Instantiate CUTLASS kernel depending on templates Gemm gemm; diff --git a/examples/75_blackwell_grouped_gemm/CMakeLists.txt b/examples/75_blackwell_grouped_gemm/CMakeLists.txt index e8a19aa4..b44659f7 100644 --- a/examples/75_blackwell_grouped_gemm/CMakeLists.txt +++ b/examples/75_blackwell_grouped_gemm/CMakeLists.txt @@ -31,16 +31,16 @@ -set(TEST_RANDOM --iterations=0) # Random problem sizes -set(TEST_RANDOM_LARGE_GROUP --groups=50 --iterations=0) # Random problem sizes +set(TEST_RANDOM --iterations=0 --raster=M --swizzle=4) # Random problem sizes +set(TEST_RANDOM_LARGE_GROUP --groups=50 --iterations=0 --raster=M --swizzle=2) # Random problem sizes set(TEST_EPILOGUE --alpha=0.5 --beta=0.5 --iterations=0) # Random problem sizes set(TEST_EPILOGUE_LARGE_GROUP --alpha=1.5 --beta=2.0 --groups=50 --iterations=0) # Random problem sizes -set(TEST_EPILOGUE_OP --beta=0.5 --iterations=1) # Random problem sizes +set(TEST_EPILOGUE_OP --beta=0.5 --iterations=1 --raster=N --swizzle=8) # Random problem sizes set(TEST_EPILOGUE_OP_LARGE_GROUP --alpha=1.5 --groups=50 --iterations=1) # Random problem sizes -set(TEST_FIXED --m=2048 --n=5120 --k=8192 --iterations=0) # Fixed problem sizes +set(TEST_FIXED --m=2048 --n=5120 --k=8192 --iterations=0 --raster=M --swizzle=8) # Fixed problem sizes set(TEST_FIXED_LARGE_GROUP --m=2048 --n=512 --k=512 --groups=51 --iterations=0) # Fixed problem sizes set(TEST_SMALL --m=256 --n=128 --iterations=0) # Small problem sizes diff --git a/examples/python/CuTeDSL/ampere/call_bypass_dlpack.py b/examples/python/CuTeDSL/ampere/call_bypass_dlpack.py index fc109f56..1a21d37a 100644 --- a/examples/python/CuTeDSL/ampere/call_bypass_dlpack.py +++ b/examples/python/CuTeDSL/ampere/call_bypass_dlpack.py @@ -29,7 +29,6 @@ import sys import os from typing import Tuple -import torch import cutlass import cutlass.cute as cute @@ -125,6 +124,8 @@ def tensor_op_gemm_wrapper( def run_tensor_op_gemm_wrapper(mnkl: Tuple[int, int, int, int]): + import torch + print("\nRunning TensorOpGemm test with:") print(f"Tensor dimensions: {mnkl}") diff --git a/examples/python/CuTeDSL/ampere/call_from_jit.py b/examples/python/CuTeDSL/ampere/call_from_jit.py index 0e1ff2c0..f4fa4339 100644 --- a/examples/python/CuTeDSL/ampere/call_from_jit.py +++ b/examples/python/CuTeDSL/ampere/call_from_jit.py @@ -60,11 +60,8 @@ import os import sys from typing import Type, Tuple -import torch - import cutlass import cutlass.cute as cute -from cutlass.torch import dtype as torch_dtype from cutlass.cute.runtime import make_ptr if __name__ == "__main__": @@ -205,6 +202,9 @@ def tensor_op_gemm_wrapper( def run_tensor_op_gemm_wrapper(mnkl: Tuple[int, int, int, int]): + import torch + from cutlass.torch import dtype as torch_dtype + print("\nRunning TensorOpGemm test with:") print(f"Tensor dimensions: {mnkl}") diff --git a/examples/python/CuTeDSL/ampere/cooperative_launch.py b/examples/python/CuTeDSL/ampere/cooperative_launch.py new file mode 100644 index 00000000..a447bf31 --- /dev/null +++ b/examples/python/CuTeDSL/ampere/cooperative_launch.py @@ -0,0 +1,627 @@ +# 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. + +""" +Cooperative Launch Example: + +This module demonstrates CUDA Cooperative Launch functionality. It implements a +global barrier that synchronizes ALL threads across the entire GPU grid. + +In traditional CUDA kernel launches, there is no guarantee that all thread blocks +will be resident on the GPU simultaneously. This means that thread blocks may +execute in waves (some finish before others start) and attempting to synchronize +across blocks can cause deadlock. + +**Cooperative Launch** solves this by guaranteeing that all thread blocks launch +atomically and simultaneously. + +For more details, see the CUDA Programming Guide official documentation: +https://docs.nvidia.com/cuda/cuda-programming-guide/04-special-topics/cooperative-groups.html#when-to-use-cudalaunchcooperativekernel + +Cooperative Launch Limitations: + +Cooperative launch has strict grid size constraints. +If you exceed this limit, cudaLaunchCooperativeKernel returns +cudaErrorCooperativeLaunchTooLarge. + +This example demonstrates both a successful cooperative launch with a small grid +and an expected failure when exceeding the grid size limit. + +Usage: + +Run directly: + $ python cooperative_launch.py + +This will: + 1. Demonstrate expected failure with too many thread blocks + 2. Successfully run a cooperative kernel with grid-wide barrier + 3. Print confirmation that all threads synchronized successfully + +""" + +from typing import List, Optional +import cutlass +import cutlass.cute as cute +from cutlass._mlir import ir +from cutlass.cutlass_dsl import ( + dsl_user_op, # Decorator for user-defined device operations + DSLCudaRuntimeError, # Exception type for CUDA runtime errors + extract_mlir_values, # Extract MLIR values from the object + new_from_mlir_values, # Create a new instance from MLIR values +) + +# Function to check cuda errors +from cutlass.base_dsl.runtime.cuda import checkCudaErrors + +# LLVM dialect for inline PTX assembly generation +from cutlass._mlir.dialects import llvm + +# CUDA Python bindings for runtime API (memory allocation, synchronization, etc.) +import cuda.bindings.runtime as cuda_runtime + + +class GlobalBarrier: + """ + A grid-wide barrier for synchronizing ALL thread blocks on the GPU. + + This class implements a cooperative barrier that enables grid-wide + synchronization. It requires cooperative launch to function correctly. + + Design Overview: + + The barrier uses a single 32-bit integer in global memory with the + following bit layout: + + ┌──────────────────────────────────────────────────────────────────┐ + │ Bit 31 │ Bits 30-0 │ + │ ────────── │ ───────────────────────────────────────────────────│ + │ Phase Bit │ Arrival Counter (supports up to 2^31 - 1 blocks) │ + └──────────────────────────────────────────────────────────────────┘ + + Capacity: + + - Maximum thread blocks: 2^31 - 1 = 2,147,483,647 blocks + + Memory Ordering: + + The barrier uses specific memory ordering semantics: + + - Release semantics on arrival (atom.add.release.gpu) + + - Acquire semantics on wait (ld.global.acquire.gpu) + + Usage Example: + + Host-side setup: + + >>> barrier_ptr = GlobalBarrier.allocate() # Allocate barrier memory + + Device-side usage (inside a kernel): + + >>> barrier = GlobalBarrier(barrier_ptr) + >>> + >>> # Do some work... + >>> + >>> barrier.arrive_and_wait() # Synchronize all blocks + >>> + >>> # All blocks proceed together after this point + + Warning: + + This barrier requires cooperative launch! Using it with a regular launch + can result in a deadlock because not all thread blocks may be resident + simultaneously. + """ + + @staticmethod + def allocate() -> cute.runtime.Pointer: + """ + Allocate and initialize barrier memory on the GPU. + + This function allocates device memory for the barrier. + It must be called before launching any kernel that uses the barrier. + """ + ptr = checkCudaErrors(cuda_runtime.cudaMalloc(4)) + + # This sets all 32 bits to 0: + # - Phase bit (bit 31) = 0 + # - Counter (bits 30-0) = 0 + checkCudaErrors(cuda_runtime.cudaMemset(ptr, 0, 4)) + + # Create a pointer with the following properties: + # - Type: Uint32 (32-bit unsigned integer) + # - Address: the allocated device pointer + # - Address Space: gmem (global memory) + barrier_ptr = cute.runtime.make_ptr( + cutlass.Uint32, # Element type + ptr, # Raw CUDA pointer + cute.AddressSpace.gmem, # Memory address space + ) + + return barrier_ptr + + @staticmethod + def free(barrier_ptr: cute.Pointer): + """ + Free the barrier memory on the GPU. + + This function frees the device memory for the barrier. + It must be called after the barrier is no longer needed. + """ + checkCudaErrors(cuda_runtime.cudaFree(barrier_ptr._pointer)) + + @dsl_user_op + def __init__( + self, + barrier_ptr: cute.Pointer, + *, + phase: Optional[cutlass.Uint32] = None, + is_leader: Optional[cutlass.Boolean] = None, + number_of_thread_blocks: Optional[cutlass.Uint32] = None, + loc=None, + ip=None, + ): + """ + Initialize a GlobalBarrier instance on the device. + + This constructor is called by each thread when the kernel + starts. It sets up the barrier state for this thread's participation + in grid-wide synchronization. + + Each thread stores the following: + + - A reference to the shared barrier memory + - Whether it's the leader thread of its block + - The current phase for barrier tracking + - The total number of thread blocks in the grid + """ + # The barrier is shared across ALL thread blocks, so it must be in + # global memory. Shared memory (smem) is block-local and wouldn't work. + if barrier_ptr.memspace != cute.AddressSpace.gmem: + raise ValueError( + "GlobalBarrier requires barrier_ptr to be in global memory (gmem)" + ) + + # Store barrier pointer reference + self.barrier_ptr = barrier_ptr + + # Initialize phase tracking + # Phase starts at 0 for the first barrier, then alternates: + # First barrier: wait for phase 1 + # Second barrier: wait for phase 0 + # Third barrier: wait for phase 1 + # ... and so on + if phase is not None: + self.phase = phase + else: + self.phase = cutlass.Uint32(0) + + if is_leader is not None: + self.is_leader = is_leader + else: + # Determine if this thread is the block leader + # Get this thread's position within its block + tidx, tidy, tidz = cute.arch.thread_idx() + + # Leader is the thread at position (0, 0, 0) in the block + # We use bitwise AND to combine the three conditions efficiently + self.is_leader = ( + cutlass.Boolean(tidx == 0) # First in X dimension + & cutlass.Boolean(tidy == 0) # First in Y dimension + & cutlass.Boolean(tidz == 0) # First in Z dimension + ) + + if number_of_thread_blocks is not None: + self.number_of_thread_blocks = number_of_thread_blocks + else: + # Calculate total number of thread blocks in the grid + # Get grid dimensions (how many blocks in each dimension) + gidx, gidy, gidz = cute.arch.grid_dim() + + # Total blocks = gridDim.x × gridDim.y × gridDim.z + # This is needed to know when ALL blocks have arrived + self.number_of_thread_blocks = cutlass.Uint32(gidx * gidy * gidz) + + @dsl_user_op + @cute.jit + def arrive(self, *, loc=None, ip=None): + """ + Arrive at the barrier without waiting. + + This signals that the calling thread block has reached the barrier + point, but does not wait for other blocks. Use this when you want + to overlap computation with barrier synchronization. + + This method must be called by ALL threads in the block, + not just the leader. The internal block-level sync ensures all + threads in the block agree before the leader signals arrival. + """ + # Ensure ALL threads in this block have reached this point before + # the leader signals arrival. This is critical for correctness! + cute.arch.sync_threads(loc=loc, ip=ip) + + # Only the leader thread performs atomic operations to minimize + # contention on the barrier memory location + if self.is_leader: + # Atomically increment the arrival counter by 1 + # The atomic add returns the value before the add, so we add 1 + # to get the current value after our arrival + barrier_value = ( + self._increment_barrier(cutlass.Uint32(1), loc=loc, ip=ip) + 1 + ) + + # Check if we're the last block to arrive + # Mask out the phase bit (bit 31) to get just the counter value + # Compare against total number of thread blocks + if (barrier_value & ~(1 << 31)) == self.number_of_thread_blocks: + # Flip phase and reset counter + # We add a value that simultaneously: + # 1. Flips bit 31 (adds 2^31) + # 2. Resets counter to 0 (subtracts N, where N was the count) + # + # Example with 8 blocks: + # Current: 0x00000008 (phase=0, counter=8) + # Add: 0x80000000 - 8 = 0x7FFFFFF8 + # Result: 0x80000000 (phase=1, counter=0) ✓ + # + # This works because we're doing modular arithmetic and the + # counter wraps correctly + self._increment_barrier( + cutlass.Uint32((1 << 31) - self.number_of_thread_blocks), + loc=loc, + ip=ip, + ) + + def _read_barrier(self, *, loc=None, ip=None) -> cutlass.Uint32: + """ + Read the barrier value with acquire memory semantics. + + This is an internal method that reads the 32-bit barrier value from + global memory using GPU-scope acquire semantics. + + Notes + ----- + PTX Instruction: + Uses ld.global.acquire.gpu.b32 which is a: + + - Global memory load (ld.global) + - With acquire semantics (.acquire) + - At GPU scope (.gpu) - visible across all thread blocks + - For 32-bit data (.b32) + + Inline Assembly: + We use LLVM inline assembly because CuTe DSL may not have a direct + high-level API for acquire loads. The assembly string format: + + - $0: Output operand (the loaded value) + - $1: Input operand (the address to load from) + + """ + # Use inline PTX assembly for the acquire-semantics load + return cutlass.Uint32( + llvm.inline_asm( + # Return type: 32-bit unsigned integer + cutlass.Uint32.mlir_type, + # Input arguments: barrier pointer address + # We convert the pointer to an integer (64-bit address) + [self.barrier_ptr.toint().ir_value(loc=loc, ip=ip)], + # PTX instruction + "ld.global.acquire.gpu.b32 $0, [$1];", + # Constraint string + # "=r" : Output is a 32-bit register (write-only) + # "l" : Input is a 64-bit register (pointer address) + "=r,l", + # Assembly attributes + # Mark as having side effects + has_side_effects=True, + # No special stack alignment needed + is_align_stack=False, + # Use AT&T syntax (required for LLVM inline asm) + asm_dialect=llvm.AsmDialect.AD_ATT, + # MLIR location and insertion point + loc=loc, + ip=ip, + ) + ) + + def _increment_barrier( + self, value: cutlass.Uint32, *, loc=None, ip=None + ) -> cutlass.Uint32: + """ + Atomically increment the barrier with release memory semantics. + + This is an internal method that performs an atomic add on the barrier + value using GPU-scope release semantics. + + Notes + ----- + PTX Instruction: + Uses atom.add.release.gpu.u32 which is a: + + - Atomic operation (atom) + - Addition (.add) + - With release semantics (.release) + - At GPU scope (.gpu) + - For unsigned 32-bit integers (.u32) + + Atomicity: + The atomic add is guaranteed to be indivisible - no other thread + can see a partial update or interleave with this operation. + + Return Value: + Atomic operations return the OLD value, not the new value. + This is why the caller adds 1 to get the current count. + """ + # Atomic add using inline PTX assembly with release semantics + return cutlass.Uint32( + llvm.inline_asm( + # Return type: 32-bit unsigned integer (the old value) + cutlass.Uint32.mlir_type, + # Input arguments: (barrier address, value to add) + [ + self.barrier_ptr.toint().ir_value( + loc=loc, ip=ip + ), # Barrier address + value.ir_value(loc=loc, ip=ip), # Value to add + ], + # PTX instruction + "atom.add.release.gpu.u32 $0, [$1], $2;", + # Constraint string + # "=r" : Output is a 32-bit register + # "l" : First input is 64-bit (pointer) + # "r" : Second input is 32-bit (value) + "=r,l,r", + # Assembly attributes + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + # MLIR metadata + loc=loc, + ip=ip, + ) + ) + + @dsl_user_op + @cute.jit + def wait(self, *, loc=None, ip=None): + """ + Wait for all thread blocks to arrive at the barrier. + + This method blocks (spins) until all thread blocks have called + arrive() on the barrier. It does NOT signal arrival itself - use + arrive_and_wait() if you need to both arrive and wait. + + IMPORTANT: This method MUST be called by ALL threads in the block. + The internal sync_threads ensures all threads proceed together. + + Algorithm: + - Leader thread spins, reading barrier with acquire semantics + - Waits until phase bit matches expected value + - Block-level sync ensures all threads proceed together + - Update local phase tracking for next barrier + + """ + # Leader thread: spin-wait for phase flip + if self.is_leader: + # Calculate expected phase (opposite of current phase) + # XOR with 1 flips: 0→1, 1→0 + expected = self.phase ^ 1 + + # Initial read of barrier value + barrier_value = self._read_barrier(loc=loc, ip=ip) + + # Spin loop: wait until phase matches expected + # Extract phase bit (bit 31) + # Compare against expected phase value + while (barrier_value >> 31) != expected: + # Keep reading barrier until phase flips + # The acquire semantics ensure memory ordering + barrier_value = self._read_barrier(loc=loc, ip=ip) + + # Block-level synchronization + # Ensure all threads in the block wait for the leader to see the + # phase flip before any thread proceeds + cute.arch.sync_threads(loc=loc, ip=ip) + + # Update phase for next barrier + # Flip local phase: 0→1 or 1→0 + # This prepares for the next barrier synchronization + self.phase = self.phase ^ 1 + + @dsl_user_op + def arrive_and_wait(self, *, loc=None, ip=None): + """ + Arrive at the barrier AND wait for all other thread blocks. + + This is the most common barrier operation - it combines arrive() + and wait() into a single call. All thread blocks will be synchronized + after this call returns. + + IMPORTANT: This method MUST be called by ALL threads in the block. + + Semantics + --------- + + Logically equivalent to: + + >>> barrier.arrive() # Signal we've reached this point + >>> barrier.wait() # Wait for everyone else + """ + # Execute both phases: arrive then wait + self.arrive(loc=loc, ip=ip) + self.wait(loc=loc, ip=ip) + + def __extract_mlir_values__(self) -> List[ir.Value]: + """ + Extract MLIR values from the GlobalBarrier instance. + """ + + assert len(extract_mlir_values(self.barrier_ptr)) == 1 + assert len(extract_mlir_values(self.is_leader)) == 1 + assert len(extract_mlir_values(self.phase)) == 1 + assert len(extract_mlir_values(self.number_of_thread_blocks)) == 1 + + return ( + extract_mlir_values(self.barrier_ptr) + + extract_mlir_values(self.is_leader) + + extract_mlir_values(self.phase) + + extract_mlir_values(self.number_of_thread_blocks) + ) + + def __new_from_mlir_values__(self, values: List[ir.Value]) -> "GlobalBarrier": + """ + Create a new GlobalBarrier instance from MLIR values. + """ + assert len(values) == 4, f"Expected 4 IR values, but got {len(values)}" + return GlobalBarrier( + barrier_ptr=new_from_mlir_values(self.barrier_ptr, [values[0]]), + is_leader=new_from_mlir_values(self.is_leader, [values[1]]), + phase=new_from_mlir_values(self.phase, [values[2]]), + number_of_thread_blocks=new_from_mlir_values( + self.number_of_thread_blocks, [values[3]] + ), + ) + + +@cute.kernel +def cooperative_kernel(barrier_ptr: cute.Pointer): + """ + Example kernel demonstrating cooperative launch with grid-wide barrier. + + This kernel shows how to use the GlobalBarrier class to synchronize all + thread blocks in a grid. It performs 10 iterations, with a barrier + synchronization after each iteration. + + Launch Requirements: This kernel MUST be launched with cooperative=True. + + """ + # Initialize the barrier for this thread + # Each thread creates its own GlobalBarrier instance, all sharing the + # same underlying barrier_ptr in global memory + barrier = GlobalBarrier(barrier_ptr=barrier_ptr) + + for i in range(10): + # Synchronize all thread blocks across the entire grid + # After this call, ALL blocks have completed iterations 0..i + barrier.arrive_and_wait() + + # Get block and thread indices + bidx, bidy, bidz = cute.arch.block_idx() + tidx, tidy, tidz = cute.arch.thread_idx() + + # Check if this is the leader block (first block in the grid) + leader_cluster = bidx == 0 and bidy == 0 and bidz == 0 + + # Check if this is the leader thread (first thread in the block) + leader_thread = tidx == 0 and tidy == 0 and tidz == 0 + + # Only the single leader thread of the leader block prints + if leader_cluster and leader_thread: + cute.printf("All threads arrived at barrier for the %dth iteration\n", i) + + +# ============================================================================= +# KERNEL LAUNCH WRAPPERS +# ============================================================================= + + +@cute.jit +def run_cooperative_kernel(barrier_ptr: cute.runtime.Pointer): + """ + Launch the cooperative kernel with a reasonable grid size. + + This wrapper launches the cooperative_kernel with a grid of 8 thread blocks + (2×2×2), where each block contains 128 threads (32×2×2). + + Notes: The cooperative=True flag is ESSENTIAL. It tells CUDA to: + - Verify the grid fits within hardware limits + - Launch all blocks atomically + """ + cooperative_kernel(barrier_ptr).launch( + grid=(2, 2, 2), # 8 thread blocks + block=(32, 2, 2), # 128 threads per block + cooperative=True, # Enable cooperative launch semantics + ) + + +@cute.jit +def xfail_run_cooperative_kernel(barrier_ptr: cute.runtime.Pointer): + """ + Demonstrate cooperative launch failure with an oversized grid. + + This wrapper intentionally launches with a grid that exceeds + the limits, demonstrating how cooperative launch fails. + + This launch is expected to fail with cudaErrorCooperativeLaunchTooLarge. + + This demonstrates proper error handling for cooperative launch. + + See Also + -------- + The main() function shows how to properly catch and handle this error. + """ + # Attempt to launch with way too many blocks + cooperative_kernel(barrier_ptr).launch( + grid=(10000, 1, 1), # 10,000 blocks + block=(1024, 1, 1), # 1,024 threads per block + cooperative=True, # Cooperative launch will reject this + ) + + +if __name__ == "__main__": + # Initialize CUDA context + cutlass.cuda.initialize_cuda_context() + + # Allocate barrier memory + # Allocate 4 bytes in device global memory for the barrier state + barrier_ptr = GlobalBarrier.allocate() + + # Demonstrate expected failure (grid too large) + expectedly_failed = False + try: + # Attempt to launch with 10,000 blocks - this WILL fail + xfail_run_cooperative_kernel(barrier_ptr) + except DSLCudaRuntimeError as e: + # Verify we got the expected error code + assert ( + e.error_code == cuda_runtime.cudaError_t.cudaErrorCooperativeLaunchTooLarge + ) + expectedly_failed = True + finally: + # Ensure the failure actually happened (test validation) + assert expectedly_failed + + # Run successful cooperative kernel + # Launch with a reasonable grid size that fits hardware constraints + run_cooperative_kernel(barrier_ptr) + + # Synchronize and clean up + checkCudaErrors(cuda_runtime.cudaDeviceSynchronize()) + + # Free the barrier memory + GlobalBarrier.free(barrier_ptr) diff --git a/examples/python/CuTeDSL/ampere/elementwise_add.py b/examples/python/CuTeDSL/ampere/elementwise_add.py index 941bda44..92fd7e25 100644 --- a/examples/python/CuTeDSL/ampere/elementwise_add.py +++ b/examples/python/CuTeDSL/ampere/elementwise_add.py @@ -28,7 +28,6 @@ import argparse -import torch import time from typing import Type @@ -36,7 +35,6 @@ from typing import Type import cutlass import cutlass.cute as cute import cutlass.cute.testing as testing -import cutlass.torch as cutlass_torch from cutlass.cute.runtime import from_dlpack """ @@ -252,6 +250,8 @@ def elementwise_add(mA, mB, mC, copy_bits: cutlass.Constexpr = 128): cC = cute.zipped_divide(idC, tiler=tiler_mn) print(f"[DSL INFO] coord tensor = {cC.type}") + kernel_name = f"cutlass_dsl_elementwise_add_kernel" + elementwise_add_kernel.set_name_prefix(kernel_name) elementwise_add_kernel(gA, gB, gC, cC, mC.shape, thr_layout, val_layout).launch( grid=[cute.size(gC, mode=[1]), 1, 1], block=[cute.size(tv_layout, mode=[0]), 1, 1], @@ -270,6 +270,12 @@ def run_elementwise_add( warmup_iterations=2, iterations=200, ): + import torch + import cutlass.torch as cutlass_torch + + if not torch.cuda.is_available(): + raise RuntimeError("Ampere GPU is required to run this example!") + print("\nRunning Elementwise Add test with:") print(f"Tensor dimensions: [{M}, {N}]") print(f"Input and Output Data type: {dtype}") @@ -304,6 +310,8 @@ def run_elementwise_add( else: c_tensor = c + elementwise_add.set_name_prefix("host_prefix") + print("Compiling kernel with cute.compile ...") start_time = time.time() compiled_func = cute.compile( @@ -386,9 +394,6 @@ if __name__ == "__main__": args = parser.parse_args() - if not torch.cuda.is_available(): - raise RuntimeError("Ampere GPU is required to run this example!") - run_elementwise_add( args.M, args.N, diff --git a/examples/python/CuTeDSL/ampere/elementwise_add_autotune.py b/examples/python/CuTeDSL/ampere/elementwise_add_autotune.py new file mode 100644 index 00000000..483c93ed --- /dev/null +++ b/examples/python/CuTeDSL/ampere/elementwise_add_autotune.py @@ -0,0 +1,372 @@ +# 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. + + +import argparse +from typing import Any, Callable, Type + +import cutlass +import cutlass.cute as cute +import cutlass.cute.testing as testing + +""" +In this example we revisit the elementwise add example and use the autotune_jit decorator to +autotune the kernel. + +To run this example: + +.. code-block:: bash + + python examples/ampere/elementwise_add_autotune.py --M 3 --N 12 + python examples/ampere/elementwise_add_autotune.py --M 1024 --N 512 + python examples/ampere/elementwise_add_autotune.py --M 1024 --N 1024 --benchmark --warmup_iterations 2 --iterations 1000 + +""" + + +@cute.kernel +def elementwise_add_kernel( + gA: cute.Tensor, + gB: cute.Tensor, + gC: cute.Tensor, + cC: cute.Tensor, # coordinate tensor + shape: cute.Shape, + thr_layout: cute.Layout, + val_layout: cute.Layout, +): + tidx, _, _ = cute.arch.thread_idx() + bidx, _, _ = cute.arch.block_idx() + + # slice for CTAs + # logical id -> address + blk_coord = ((None, None), bidx) + blkA = gA[blk_coord] # (TileM,TileN) + blkB = gB[blk_coord] # (TileM,TileN) + blkC = gC[blk_coord] # (TileM,TileN) + blkCrd = cC[blk_coord] # (TileM, TileN) + + # # declare the atoms which will be used later for memory copy + copy_atom_load = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), gA.element_type) + copy_atom_store = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), gC.element_type) + + tiled_copy_A = cute.make_tiled_copy_tv(copy_atom_load, thr_layout, val_layout) + tiled_copy_B = cute.make_tiled_copy_tv(copy_atom_load, thr_layout, val_layout) + tiled_copy_C = cute.make_tiled_copy_tv(copy_atom_store, thr_layout, val_layout) + + thr_copy_A = tiled_copy_A.get_slice(tidx) + thr_copy_B = tiled_copy_B.get_slice(tidx) + thr_copy_C = tiled_copy_C.get_slice(tidx) + + thrA = thr_copy_A.partition_S(blkA) + thrB = thr_copy_B.partition_S(blkB) + thrC = thr_copy_C.partition_S(blkC) + + # allocate fragments for gmem->rmem + frgA = cute.make_rmem_tensor_like(thrA) + frgB = cute.make_rmem_tensor_like(thrB) + frgC = cute.make_rmem_tensor_like(thrC) + + thrCrd = thr_copy_C.partition_S(blkCrd) + frgPred = cute.make_rmem_tensor(thrCrd.shape, cutlass.Boolean) + + for i in range(0, cute.size(frgPred), 1): + val = cute.elem_less(thrCrd[i], shape) + frgPred[i] = val + + # Print per thread predicate mask + # if tidx == 0 and bidx == 0: + # cute.printf("block_dim = {}", cute.arch.grid_dim()) + # cute.printf("shape = {}", shape) + # cute.print_tensor(thrA) + # cute.print_tensor(thrB) + # cute.print_tensor(frgPred) + + ########################################################## + # Move data to reg address space + ########################################################## + + cute.copy(copy_atom_load, thrA, frgA, pred=frgPred) + cute.copy(copy_atom_load, thrB, frgB, pred=frgPred) + + # if tidx == 0 and bidx == 0: + # cute.print_tensor(frgA) + # cute.print_tensor(frgB) + + # Load data before use. The compiler will optimize the copy and load + # operations to convert some memory ld/st into register uses. + result = frgA.load() + frgB.load() + + # Save the results back to registers. Here we reuse b's registers. + frgC.store(result) + + # Copy the results back to c + cute.copy(copy_atom_store, frgC, thrC, pred=frgPred) + + +@testing.autotune_jit( + params_dict={"copy_bits": [64, 128]}, + update_on_change=["M", "N"], + warmup_iterations=100, + iterations=100, +) +@cute.jit +def elementwise_add_autotune(mA, mB, mC, M, N, copy_bits: cutlass.Constexpr = 128): + dtype = mA.element_type + vector_size = copy_bits // dtype.width + + thr_layout = cute.make_ordered_layout((4, 32), order=(1, 0)) + val_layout = cute.make_ordered_layout((4, vector_size), order=(1, 0)) + tiler_mn, tv_layout = cute.make_layout_tv(thr_layout, val_layout) + + gA = cute.zipped_divide(mA, tiler_mn) # ((TileM,TileN),(RestM,RestN)) + gB = cute.zipped_divide(mB, tiler_mn) # ((TileM,TileN),(RestM,RestN)) + gC = cute.zipped_divide(mC, tiler_mn) # ((TileM,TileN),(RestM,RestN)) + idC = cute.make_identity_tensor(mC.shape) + cC = cute.zipped_divide(idC, tiler=tiler_mn) + + elementwise_add_kernel(gA, gB, gC, cC, mC.shape, thr_layout, val_layout).launch( + grid=[cute.size(gC, mode=[1]), 1, 1], + block=[cute.size(tv_layout, mode=[0]), 1, 1], + ) + + +class ElementwiseAddWrapper: + """ + This class mimics more advanced kernel development, where a class encapsulates + pieces of the kernel implementation. + + The can_implement method can be used to check if the kernel can be implemented + for the given arguments. + + The __call__ method is the actual cute.jit function. + + """ + + def __init__(self, copy_bits: cutlass.Constexpr = 128): + self.copy_bits = copy_bits + + def can_implement(self, mA, mB, mC, M, N): + return self.copy_bits in [64, 128] + + @cute.jit + def __call__(self, mA, mB, mC, M, N): + dtype = mA.element_type + vector_size = self.copy_bits // dtype.width + + thr_layout = cute.make_ordered_layout((4, 32), order=(1, 0)) + val_layout = cute.make_ordered_layout((4, vector_size), order=(1, 0)) + tiler_mn, tv_layout = cute.make_layout_tv(thr_layout, val_layout) + + gA = cute.zipped_divide(mA, tiler_mn) # ((TileM,TileN),(RestM,RestN)) + gB = cute.zipped_divide(mB, tiler_mn) # ((TileM,TileN),(RestM,RestN)) + gC = cute.zipped_divide(mC, tiler_mn) # ((TileM,TileN),(RestM,RestN)) + idC = cute.make_identity_tensor(mC.shape) + cC = cute.zipped_divide(idC, tiler=tiler_mn) + + elementwise_add_kernel(gA, gB, gC, cC, mC.shape, thr_layout, val_layout).launch( + grid=[cute.size(gC, mode=[1]), 1, 1], + block=[cute.size(tv_layout, mode=[0]), 1, 1], + ) + + +def tune_class(mA, mB, mC, M, N): + """ + This function is used to autotune the elementwise add kernel which is wrapped in a class. + An internal function is defined to compile the class with the given arguments. + The internal function is then passed to the benchmarking.tune function to autotune. + The best parameters are then used to instantiate the class. + + :param mA: Input tensor A + :type mA: cute.Tensor + :param mB: Input tensor B + :type mB: cute.Tensor + :param mC: Output tensor C + :type mC: cute.Tensor + :param M: Number of rows in the input tensors + :type M: int + :param N: Number of columns in the input tensors + :type N: int + :return: An instance of the ElementwiseAddWrapper class with the best parameters + :rtype: ElementwiseAddWrapper + """ + + def compile_class(a, b, c, M, N, copy_bits=128) -> Callable[[], Any]: + kernel = ElementwiseAddWrapper(copy_bits) + if not kernel.can_implement(a, b, c, M, N): + raise ValueError(f"Cannot implement kernel for copy_bits={copy_bits}") + compiled_kernel = cute.compile(kernel, a, b, c, M, N) + return lambda: compiled_kernel(a, b, c, M, N) + + params = testing.tune( + compile_class, + params_dict={"copy_bits": [1, 64, 128]}, + kernel_arguments=testing.JitArguments(mA, mB, mC, M, N), + ) + return ElementwiseAddWrapper(**params) + + +def run_elementwise_add( + M_start, + M_range, + M_step, + N_start, + N_range, + N_step, + dtype: Type[cutlass.Numeric], + skip_ref_check=False, + warmup_iterations=2, + iterations=200, +): + import torch + import cutlass.torch as cutlass_torch + + if not torch.cuda.is_available(): + raise RuntimeError("Ampere GPU is required to run this example!") + + for M in range(M_start, M_start + M_range + 1, M_step): + for N in range(N_start, N_start + N_range + 1, N_step): + print("\nRunning Elementwise Add test with:") + print(f"Tensor dimensions: [{M}, {N}]") + print(f"Input and Output Data type: {dtype}") + + torch_dtype = cutlass_torch.dtype(dtype) + if dtype.is_integer: + a = torch.randint( + 0, 10, (M, N), device=torch.device("cuda"), dtype=torch_dtype + ) + b = torch.randint( + 0, 10, (M, N), device=torch.device("cuda"), dtype=torch_dtype + ) + else: + a = torch.randn(M, N, device=torch.device("cuda"), dtype=torch_dtype) + b = torch.randn(M, N, device=torch.device("cuda"), dtype=torch_dtype) + + c = torch.zeros_like(a) + + print("Input tensor shapes:") + print(f"a: {a.shape}, dtype: {a.dtype}") + print(f"b: {b.shape}, dtype: {b.dtype}") + print(f"c: {c.shape}, dtype: {c.dtype}\n") + + elementwise_class = tune_class(a, b, c, M, N) + + if not skip_ref_check: + print("Verifying results for class ...") + torch.testing.assert_close(a + b, c) + print("Results verified successfully!") + c = torch.zeros_like(a) + + elementwise_add_autotune(a, b, c, M, N) + + if not skip_ref_check: + print("Verifying results for autotuned function ...") + torch.testing.assert_close(a + b, c) + print("Results verified successfully!") + + def generate_kernel_arguments(): + if dtype.is_integer: + a = torch.randint( + 0, 10, (M, N), device=torch.device("cuda"), dtype=torch_dtype + ) + b = torch.randint( + 0, 10, (M, N), device=torch.device("cuda"), dtype=torch_dtype + ) + else: + a = torch.randn( + M, N, device=torch.device("cuda"), dtype=torch_dtype + ) + b = torch.randn( + M, N, device=torch.device("cuda"), dtype=torch_dtype + ) + + c = torch.zeros_like(a) + + return testing.JitArguments(a, b, c, M, N) + + avg_time_us = testing.benchmark( + elementwise_add_autotune, + workspace_generator=generate_kernel_arguments, + workspace_count=10, + warmup_iterations=warmup_iterations, + iterations=iterations, + ) + + # Print execution results + print( + f"Kernel execution time for cute.jit kernel with M={M}, N={N}: {avg_time_us / 1e3:.4f} ms" + ) + print( + f"Achieved memory throughput for M={M}, N={N}: {(3 * a.numel() * dtype.width // 8) / (avg_time_us / 1e6) / 1e9:.2f} GB/s" + ) + + compiled_class = cute.compile(elementwise_class, a, b, c, M, N) + + avg_time_us = testing.benchmark( + compiled_class, + workspace_generator=generate_kernel_arguments, + workspace_count=10, + warmup_iterations=warmup_iterations, + iterations=iterations, + ) + + print( + f"Kernel execution time for Class Wrapper with M={M}, N={N}: {avg_time_us / 1e3:.4f} ms" + ) + print( + f"Achieved memory throughput for M={M}, N={N}: {(3 * a.numel() * dtype.width // 8) / (avg_time_us / 1e6) / 1e9:.2f} GB/s" + ) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + description="example of elementwise add to demonstrate the numpy/pytorch as input for kernels" + ) + parser.add_argument("--M", default=1024, type=int) + parser.add_argument("--M_range", default=0, type=int) + parser.add_argument("--M_step", default=1024, type=int) + parser.add_argument("--N", default=1024, type=int) + parser.add_argument("--N_range", default=0, type=int) + parser.add_argument("--N_step", default=1024, type=int) + parser.add_argument("--warmup_iterations", default=2, type=int) + parser.add_argument("--iterations", default=100, type=int) + parser.add_argument("--skip_ref_check", action="store_true") + + args = parser.parse_args() + run_elementwise_add( + args.M, + args.M_range, + args.M_step, + args.N, + args.N_range, + args.N_step, + dtype=cutlass.Float32, + skip_ref_check=args.skip_ref_check, + warmup_iterations=args.warmup_iterations, + iterations=args.iterations, + ) + print("\nPASS") diff --git a/examples/python/CuTeDSL/ampere/elementwise_apply.py b/examples/python/CuTeDSL/ampere/elementwise_apply.py index 12f93df8..2d39cbce 100644 --- a/examples/python/CuTeDSL/ampere/elementwise_apply.py +++ b/examples/python/CuTeDSL/ampere/elementwise_apply.py @@ -36,8 +36,6 @@ from typing import List, Type import cuda.bindings.driver as cuda import cutlass.cute as cute import cutlass.cute.testing as testing -import cutlass.torch as cutlass_torch -import torch from cutlass.cute.runtime import from_dlpack import cutlass @@ -274,6 +272,8 @@ def leaky_relu(x, alpha, *, loc=None, ip=None): def leaky_relu_ref(x, alpha): + import torch + return torch.where(x > 0, x, alpha * x) @@ -287,6 +287,9 @@ def run_and_verify( warmup_iterations=2, iterations=100, ): + import torch + import cutlass.torch as cutlass_torch + if not torch.cuda.is_available(): raise RuntimeError("NVIDIA GPU is required to run this example!") diff --git a/examples/python/CuTeDSL/ampere/flash_attention_v2.py b/examples/python/CuTeDSL/ampere/flash_attention_v2.py index 5f936e72..95e3497b 100644 --- a/examples/python/CuTeDSL/ampere/flash_attention_v2.py +++ b/examples/python/CuTeDSL/ampere/flash_attention_v2.py @@ -30,13 +30,11 @@ import argparse from types import SimpleNamespace from typing import Type, Callable -import torch import cuda.bindings.driver as cuda import cutlass.cute.testing as testing import cutlass import cutlass.cute as cute from cutlass.cute.nvgpu import cpasync, warp -import cutlass.torch as cutlass_torch from cutlass.cute.runtime import from_dlpack import cutlass.pipeline as pipeline import cutlass.utils as utils @@ -1162,6 +1160,9 @@ def run( use_cold_l2: bool = False, **kwargs, ): + import torch + import cutlass.torch as cutlass_torch + # Skip unsupported testcase if not FlashAttentionForwardAmpere.can_implement( dtype, @@ -1237,8 +1238,12 @@ def run( torch_stream = torch.cuda.current_stream() # Get the raw stream pointer as a CUstream current_stream = cuda.CUstream(torch_stream.cuda_stream) + # Pass compile options if needed + compile_options = "" # compile the fa2 forward pass - compiled_fa2_fwd = cute.compile(fa2_fwd, q, k, v, o, softmax_scale, current_stream) + compiled_fa2_fwd = cute.compile( + fa2_fwd, q, k, v, o, softmax_scale, current_stream, options=compile_options + ) if not skip_ref_check: compiled_fa2_fwd(q, k, v, o, softmax_scale, current_stream) @@ -1317,7 +1322,6 @@ if __name__ == "__main__": default=False, help="Use circular buffer tensor sets to ensure L2 cold cache", ) - args = parser.parse_args() run( args.dtype, diff --git a/examples/python/CuTeDSL/ampere/hstu_attention.py b/examples/python/CuTeDSL/ampere/hstu_attention.py index b32600e8..3b537792 100644 --- a/examples/python/CuTeDSL/ampere/hstu_attention.py +++ b/examples/python/CuTeDSL/ampere/hstu_attention.py @@ -29,10 +29,8 @@ from typing import Type import argparse -import torch import cuda.bindings.driver as cuda import cutlass -import cutlass.torch as cutlass_torch import cutlass.cute as cute from cutlass.cute.runtime import from_dlpack from cutlass._mlir.dialects import llvm @@ -838,28 +836,25 @@ class HSTUAttentionForwardAmpere(object): def run_pytorch_hstu_test( - dtype: torch.dtype, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - rab: torch.Tensor, + dtype, + q, + k, + v, + rab, is_causal: bool, ): """Generate the reference output of the HSTU attention with Pytorch. :param dtype: data type of the input tensors - :type dtype: torch.dtype :param q: query tensor - :type q: torch.Tensor :param k: key tensor - :type k: torch.Tensor :param v: value tensor - :type v: torch.Tensor :param rab: RAB tensor - :type rab: torch.Tensor :param is_causal: whether to use causal masking :type is_causal: bool """ + import torch + q = q.to(dtype) k = k.to(dtype) v = v.to(dtype) @@ -922,6 +917,9 @@ def run( """ assert dtype == cutlass.Float16 or dtype == cutlass.BFloat16 + import torch + import cutlass.torch as cutlass_torch + torch_stream = torch.cuda.current_stream() stream = cuda.CUstream(torch_stream.cuda_stream) diff --git a/examples/python/CuTeDSL/ampere/inline_ptx.py b/examples/python/CuTeDSL/ampere/inline_ptx.py index 73d54dee..0eeedeb3 100644 --- a/examples/python/CuTeDSL/ampere/inline_ptx.py +++ b/examples/python/CuTeDSL/ampere/inline_ptx.py @@ -29,8 +29,6 @@ from functools import partial from typing import Union -import torch - import cutlass.cute as cute from cutlass.cute.runtime import from_dlpack from cutlass._mlir.dialects import llvm @@ -49,7 +47,7 @@ Situations like: motivate developers to inline PTX themselves. In this example, we inline the vote.sync.ballot.b32, vote.sync.any.pred, vote.sync.all.pred, -vote.sync.uni.pred, and use the corresponding ops in nvvm_wrappers.py for the test. +vote.sync.uni.pred, and use the corresponding ops in nvvm dialect for the test. You can refer to the documentation of `inline_asm op in llvm dialect `_ and `vote.sync `_ @@ -61,8 +59,8 @@ To run this example: python examples/ampere/inline_ptx.py -The example will run the vote kernel with inline PTX and nvvm dialect separately. -The results from inline PTX and nvvm dialect will be verified correspondingly. +The example will run the vote kernel with inline ptx and nvvm dialect separately. +The results from inline ptx and nvvm dialect will be verified correspondingly. """ @@ -184,6 +182,8 @@ def vote( def run(): + import torch + ballot_ptx = torch.randint( 0, 100, (WARP_SIZE,), device=torch.device("cuda"), dtype=torch.int32 ) @@ -230,14 +230,11 @@ def run(): torch.testing.assert_close(ballot_ptx, ballot_nvvm) print("Verifying any results...") torch.testing.assert_close(any_ptx, any_nvvm) - print(torch.all(any_ptx == any(i < 10 for i in range(WARP_SIZE)))) - assert torch.all(any_ptx == any(i < 10 for i in range(WARP_SIZE))) print("Verifying all results...") torch.testing.assert_close(all_ptx, all_nvvm) - assert torch.all(all_ptx == all(i < 10 for i in range(WARP_SIZE))) print("Verifying uni results...") torch.testing.assert_close(uni_ptx, uni_nvvm) - assert torch.all(uni_ptx == (len(set(i < 10 for i in range(WARP_SIZE))) == 1)) + print("Results verified successfully!") diff --git a/examples/python/CuTeDSL/ampere/sgemm.py b/examples/python/CuTeDSL/ampere/sgemm.py index 17835dd0..3455e9b7 100644 --- a/examples/python/CuTeDSL/ampere/sgemm.py +++ b/examples/python/CuTeDSL/ampere/sgemm.py @@ -31,7 +31,6 @@ import time from typing import Tuple import cuda.bindings.driver as cuda -import torch import cutlass import cutlass.cute as cute @@ -643,6 +642,8 @@ def run( use_cold_l2: bool = False, **kwargs, ): + import torch + """Execute SIMT GEMM operation and benchmark performance. :param mnk: GEMM problem size (M, N, K, L) @@ -666,6 +667,7 @@ def run( :return: Execution time of the GEMM kernel in microseconds :rtype: float """ + torch.manual_seed(1024) print("Running Ampere SIMT GEMM example:") print(f"mnk: {mnk}") print(f"A major: {a_major}, B major: {b_major}, C major: {c_major}") @@ -851,8 +853,6 @@ if __name__ == "__main__": args = parser.parse_args() print("Running SIMT GEMM example:") - torch.manual_seed(1024) - run( args.mnk, args.a_major, diff --git a/examples/python/CuTeDSL/ampere/smem_allocator.py b/examples/python/CuTeDSL/ampere/smem_allocator.py index 3728db43..ea004ead 100644 --- a/examples/python/CuTeDSL/ampere/smem_allocator.py +++ b/examples/python/CuTeDSL/ampere/smem_allocator.py @@ -28,7 +28,6 @@ import cutlass.cute as cute import cutlass -import torch import numpy as np from cutlass.cute.runtime import from_dlpack @@ -175,6 +174,8 @@ def host( def run_and_verify(const_a, const_b, const_c): + import torch + dst_a = torch.zeros((8, 4), dtype=torch.float32, device="cuda") dst_b = torch.zeros((8, 2), dtype=torch.float32, device="cuda") dst_c = torch.zeros((16, 2), dtype=torch.float32, device="cuda") diff --git a/examples/python/CuTeDSL/ampere/tensorop_gemm.py b/examples/python/CuTeDSL/ampere/tensorop_gemm.py index ed9a455c..ad3fe0f0 100644 --- a/examples/python/CuTeDSL/ampere/tensorop_gemm.py +++ b/examples/python/CuTeDSL/ampere/tensorop_gemm.py @@ -30,12 +30,9 @@ import argparse import math from typing import Tuple, Type -import torch - import cutlass import cutlass.cute as cute import cutlass.cute.testing as testing -import cutlass.torch as cutlass_torch import cutlass.utils as utils from cutlass.cute.runtime import from_dlpack @@ -849,6 +846,9 @@ def run( use_cold_l2: bool = False, **kwargs, ): + import torch + import cutlass.torch as cutlass_torch + print("Running Ampere tensor core GEMM example:") print(f"mnkl: {mnkl}") print( diff --git a/examples/python/CuTeDSL/blackwell/blockwise_gemm/blockwise_gemm.py b/examples/python/CuTeDSL/blackwell/blockwise_gemm/blockwise_gemm.py index c6f7b5e0..a54e0fb0 100644 --- a/examples/python/CuTeDSL/blackwell/blockwise_gemm/blockwise_gemm.py +++ b/examples/python/CuTeDSL/blackwell/blockwise_gemm/blockwise_gemm.py @@ -30,13 +30,11 @@ import argparse from typing import Type, Tuple, Union import cuda.bindings.driver as cuda -import torch import cutlass import cutlass.cute as cute import cutlass.cute.testing as testing from cutlass.cute.nvgpu import cpasync, tcgen05 -import cutlass.torch as cutlass_torch import cutlass.utils as utils import cutlass.pipeline as pipeline from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait @@ -64,7 +62,7 @@ This GEMM kernel supports the following features: This GEMM works as follows: 1. DMA warp: Load A and B matrices from global memory (GMEM) to shared memory (SMEM) using TMA operations. -2. SCALE warp: Load scaleA and scaleB matrices from global memory (GMEM) to shared memory (SMEM) using non-TMA operations. +2. SCALE warp: Load scaleA and scaleB matrices from global memory (GMEM) to shared memory (SMEM) using async copy operations. 2. MMA warp: Perform matrix multiply-accumulate (MMA) operations using tcgen05.mma instruction. 3. EPILOGUE warp: - Load completed accumulator from tensor memory (TMEM) to registers (RMEM) using tcgen05.ld. @@ -1008,7 +1006,10 @@ class BlockwiseGemmKernel: ) # fence view async shared - cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) self.sched_sync_barrier.arrive_and_wait() # commit tile info pipeline tile_info_pipeline.producer_commit(tile_info_producer_state) @@ -1123,7 +1124,10 @@ class BlockwiseGemmKernel: for idx in cutlass.range(4, unroll_full=True): tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] is_valid_tile = tile_info[3] == 1 - cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) tile_info_pipeline.consumer_release(tile_info_consumer_state) tile_info_consumer_state.advance() @@ -1295,7 +1299,10 @@ class BlockwiseGemmKernel: for idx in cutlass.range(4, unroll_full=True): tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] is_valid_tile = tile_info[3] == 1 - cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) tile_info_pipeline.consumer_release(tile_info_consumer_state) tile_info_consumer_state.advance() @@ -1450,7 +1457,10 @@ class BlockwiseGemmKernel: for idx in cutlass.range(4, unroll_full=True): tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] is_valid_tile = tile_info[3] == 1 - cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) tile_info_pipeline.consumer_release(tile_info_consumer_state) tile_info_consumer_state.advance() @@ -1684,7 +1694,10 @@ class BlockwiseGemmKernel: for idx in cutlass.range(4, unroll_full=True): tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] is_valid_tile = tile_info[3] == 1 - cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) tile_info_pipeline.consumer_release(tile_info_consumer_state) tile_info_consumer_state.advance() @@ -1851,7 +1864,10 @@ class BlockwiseGemmKernel: tRS_sC[(None, None, None, c_buffer)], ) # Fence and barrier to make sure shared memory store is visible to TMA store - cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) self.epilog_sync_barrier.arrive_and_wait() # @@ -1881,7 +1897,10 @@ class BlockwiseGemmKernel: for idx in cutlass.range(4, unroll_full=True): tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] is_valid_tile = tile_info[3] == 1 - cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) tile_info_pipeline.consumer_release(tile_info_consumer_state) tile_info_consumer_state.advance() @@ -2553,6 +2572,9 @@ class BlockwiseGemmKernel: def create_tensors( l, m, n, k, a_major, b_major, cd_major, ab_dtype, c_dtype, scale_dtype ): + import torch + import cutlass.torch as cutlass_torch + torch.manual_seed(1111) a_torch_cpu = cutlass_torch.matrix(l, m, k, a_major == "m", ab_dtype) @@ -2613,6 +2635,9 @@ def run( use_cold_l2: bool = False, **kwargs, ): + import torch + import cutlass.torch as cutlass_torch + """ Prepare A/B/C tensors, launch GPU kernel, and reference checking. """ @@ -2688,6 +2713,7 @@ def run( # try to check CUDA version to decide the opt level try: from cutlass import CUDA_VERSION + opt_level = ( 3 if CUDA_VERSION.major < 13 diff --git a/examples/python/CuTeDSL/blackwell/blockwise_gemm/contiguous_grouped_gemm.py b/examples/python/CuTeDSL/blackwell/blockwise_gemm/contiguous_grouped_gemm.py index 233dda4a..bc99a15c 100644 --- a/examples/python/CuTeDSL/blackwell/blockwise_gemm/contiguous_grouped_gemm.py +++ b/examples/python/CuTeDSL/blackwell/blockwise_gemm/contiguous_grouped_gemm.py @@ -30,13 +30,11 @@ import argparse from typing import Type, Tuple, Union import cuda.bindings.driver as cuda -import torch import cutlass import cutlass.cute as cute import cutlass.cute.testing as testing from cutlass.cute.nvgpu import cpasync, tcgen05 -import cutlass.torch as cutlass_torch import cutlass.utils as utils import cutlass.pipeline as pipeline from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait @@ -80,7 +78,7 @@ This GEMM kernel supports the following features: This GEMM works as follows: 1. DMA warp: Load A and B matrices from global memory (GMEM) to shared memory (SMEM) using TMA operations. -2. SCALE warp: Load scaleA and scaleB matrices from global memory (GMEM) to shared memory (SMEM) using non-TMA operations. +2. SCALE warp: Load scaleA and scaleB matrices from global memory (GMEM) to shared memory (SMEM) using async copy operations. 2. MMA warp: Perform matrix multiply-accumulate (MMA) operations using tcgen05.mma instruction. 3. EPILOGUE warp: - Load completed accumulator from tensor memory (TMEM) to registers (RMEM) using tcgen05.ld. @@ -1034,7 +1032,10 @@ class BlockwiseContiguousGroupedGemmKernel: ) # fence view async shared - cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) self.sched_sync_barrier.arrive_and_wait() # commit tile info pipeline tile_info_pipeline.producer_commit(tile_info_producer_state) @@ -1150,7 +1151,10 @@ class BlockwiseContiguousGroupedGemmKernel: for idx in cutlass.range(4, unroll_full=True): tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] is_valid_tile = tile_info[3] == 1 - cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) tile_info_pipeline.consumer_release(tile_info_consumer_state) tile_info_consumer_state.advance() @@ -1322,7 +1326,10 @@ class BlockwiseContiguousGroupedGemmKernel: for idx in cutlass.range(4, unroll_full=True): tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] is_valid_tile = tile_info[3] == 1 - cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) tile_info_pipeline.consumer_release(tile_info_consumer_state) tile_info_consumer_state.advance() @@ -1479,7 +1486,10 @@ class BlockwiseContiguousGroupedGemmKernel: for idx in cutlass.range(4, unroll_full=True): tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] is_valid_tile = tile_info[3] == 1 - cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) tile_info_pipeline.consumer_release(tile_info_consumer_state) tile_info_consumer_state.advance() @@ -1715,7 +1725,10 @@ class BlockwiseContiguousGroupedGemmKernel: for idx in cutlass.range(4, unroll_full=True): tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] is_valid_tile = tile_info[3] == 1 - cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) tile_info_pipeline.consumer_release(tile_info_consumer_state) tile_info_consumer_state.advance() @@ -1884,7 +1897,10 @@ class BlockwiseContiguousGroupedGemmKernel: tRS_sC[(None, None, None, c_buffer)], ) # Fence and barrier to make sure shared memory store is visible to TMA store - cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) self.epilog_sync_barrier.arrive_and_wait() # @@ -1914,7 +1930,10 @@ class BlockwiseContiguousGroupedGemmKernel: for idx in cutlass.range(4, unroll_full=True): tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] is_valid_tile = tile_info[3] == 1 - cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) tile_info_pipeline.consumer_release(tile_info_consumer_state) tile_info_consumer_state.advance() @@ -2595,6 +2614,8 @@ class BlockwiseContiguousGroupedGemmKernel: def create_mask(num_groups, expect_m, fixed_m=False, m_aligned=128): + import torch + valid_m = 0 group_m_list = [] gidx_mapping = [] @@ -2632,6 +2653,9 @@ def create_tensors( scale_dtype, fixed_m=False, ): + import torch + import cutlass.torch as cutlass_torch + torch.manual_seed(1111) valid_m, group_m_list, _gidx_mapping = create_mask(l, m, fixed_m) @@ -2702,6 +2726,9 @@ def run( fixed_m: bool = False, **kwargs, ): + import torch + import cutlass.torch as cutlass_torch + """ Prepare A/B/C tensors, launch GPU kernel, and reference checking. """ diff --git a/examples/python/CuTeDSL/blackwell/blockwise_gemm/masked_grouped_gemm.py b/examples/python/CuTeDSL/blackwell/blockwise_gemm/masked_grouped_gemm.py index a0b655ae..316a5e9a 100644 --- a/examples/python/CuTeDSL/blackwell/blockwise_gemm/masked_grouped_gemm.py +++ b/examples/python/CuTeDSL/blackwell/blockwise_gemm/masked_grouped_gemm.py @@ -30,13 +30,11 @@ import argparse from typing import Type, Tuple, Union import cuda.bindings.driver as cuda -import torch import cutlass import cutlass.cute as cute import cutlass.cute.testing as testing from cutlass.cute.nvgpu import cpasync, tcgen05 -import cutlass.torch as cutlass_torch import cutlass.utils as utils import cutlass.pipeline as pipeline from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait @@ -79,7 +77,7 @@ Matrix A/C Memory Layout Diagrams: This GEMM works as follows: 1. DMA warp: Load A and B matrices from global memory (GMEM) to shared memory (SMEM) using TMA operations. -2. SCALE warp: Load scaleA and scaleB matrices from global memory (GMEM) to shared memory (SMEM) using non-TMA operations. +2. SCALE warp: Load scaleA and scaleB matrices from global memory (GMEM) to shared memory (SMEM) using async copy operations. 2. MMA warp: Perform matrix multiply-accumulate (MMA) operations using tcgen05.mma instruction. 3. EPILOGUE warp: - Load completed accumulator from tensor memory (TMEM) to registers (RMEM) using tcgen05.ld. @@ -1041,7 +1039,10 @@ class BlockwiseMaskedGroupedGemmKernel: ) # fence view async shared - cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) self.sched_sync_barrier.arrive_and_wait() # commit tile info pipeline tile_info_pipeline.producer_commit(tile_info_producer_state) @@ -1156,7 +1157,10 @@ class BlockwiseMaskedGroupedGemmKernel: for idx in cutlass.range(4, unroll_full=True): tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] is_valid_tile = tile_info[3] == 1 - cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) tile_info_pipeline.consumer_release(tile_info_consumer_state) tile_info_consumer_state.advance() @@ -1328,7 +1332,10 @@ class BlockwiseMaskedGroupedGemmKernel: for idx in cutlass.range(4, unroll_full=True): tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] is_valid_tile = tile_info[3] == 1 - cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) tile_info_pipeline.consumer_release(tile_info_consumer_state) tile_info_consumer_state.advance() @@ -1483,7 +1490,10 @@ class BlockwiseMaskedGroupedGemmKernel: for idx in cutlass.range(4, unroll_full=True): tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] is_valid_tile = tile_info[3] == 1 - cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) tile_info_pipeline.consumer_release(tile_info_consumer_state) tile_info_consumer_state.advance() @@ -1717,7 +1727,10 @@ class BlockwiseMaskedGroupedGemmKernel: for idx in cutlass.range(4, unroll_full=True): tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] is_valid_tile = tile_info[3] == 1 - cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) tile_info_pipeline.consumer_release(tile_info_consumer_state) tile_info_consumer_state.advance() @@ -1884,7 +1897,10 @@ class BlockwiseMaskedGroupedGemmKernel: tRS_sC[(None, None, None, c_buffer)], ) # Fence and barrier to make sure shared memory store is visible to TMA store - cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) self.epilog_sync_barrier.arrive_and_wait() # @@ -1914,7 +1930,10 @@ class BlockwiseMaskedGroupedGemmKernel: for idx in cutlass.range(4, unroll_full=True): tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] is_valid_tile = tile_info[3] == 1 - cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) tile_info_pipeline.consumer_release(tile_info_consumer_state) tile_info_consumer_state.advance() @@ -2586,6 +2605,8 @@ class BlockwiseMaskedGroupedGemmKernel: def create_mask(num_groups: int, m: int, fixed_m=False, tile_m=128): + import torch + # align with block_m (or block_n if swapAB) masked_m_candidates = list( filter( @@ -2617,6 +2638,9 @@ def create_tensors( scale_dtype, fixed_m=False, ): + import torch + import cutlass.torch as cutlass_torch + torch.manual_seed(1111) _gidx_mapping, masked_m = create_mask(l, m, fixed_m) @@ -2684,6 +2708,9 @@ def run( fixed_m: bool = False, **kwargs, ): + import torch + import cutlass.torch as cutlass_torch + """ Prepare A/B/C tensors, launch GPU kernel, and reference checking. """ diff --git a/examples/python/CuTeDSL/blackwell/dense_blockscaled_gemm_persistent.py b/examples/python/CuTeDSL/blackwell/dense_blockscaled_gemm_persistent.py index f3573fa7..b74d35dd 100644 --- a/examples/python/CuTeDSL/blackwell/dense_blockscaled_gemm_persistent.py +++ b/examples/python/CuTeDSL/blackwell/dense_blockscaled_gemm_persistent.py @@ -27,7 +27,7 @@ # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. import argparse -from typing import Type, Tuple, Union +from typing import Type, Tuple, Union, Literal import cuda.bindings.driver as cuda import torch @@ -41,7 +41,7 @@ import cutlass.pipeline as pipeline from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait import cutlass.utils.blackwell_helpers as sm100_utils import cutlass.utils.blockscaled_layout as blockscaled_utils -from cutlass.cute.runtime import from_dlpack +from cutlass.cute.runtime import make_ptr """ This example provides an experimental implementation of the SM100 batched dense blockscaled GEMM kernel, please note that the APIs and implementation details related to this kernel may change in future releases. @@ -117,10 +117,6 @@ Constraints: """ -def ceil_div(a, b): - return (a + b - 1) // b - - class Sm100BlockScaledPersistentDenseGemmKernel: """This class implements batched matrix multiplication (C = A x SFA x B x SFB) with support for various data types and architectural features specific to Blackwell GPUs with persistent tile scheduling and warp specialization. @@ -224,8 +220,7 @@ class Sm100BlockScaledPersistentDenseGemmKernel: num_threads=self.threads_per_warp * len((self.mma_warp_id, *self.epilog_warp_id)), ) self.smem_capacity = utils.get_smem_capacity_in_bytes("sm_100") - SM100_TMEM_CAPACITY_COLUMNS = 512 - self.num_tmem_alloc_cols = SM100_TMEM_CAPACITY_COLUMNS + self.num_tmem_alloc_cols = cute.arch.get_max_tmem_alloc_cols("sm_100") def _setup_attributes(self): """Set up configurations that are dependent on GEMM inputs @@ -375,24 +370,36 @@ class Sm100BlockScaledPersistentDenseGemmKernel: # Compute number of TMEM columns for SFA/SFB/Accumulator sf_atom_mn = 32 - self.num_sfa_tmem_cols = (self.cta_tile_shape_mnk[0] // sf_atom_mn) * mma_inst_tile_k - self.num_sfb_tmem_cols = (self.cta_tile_shape_mnk_sfb[1] // sf_atom_mn) * mma_inst_tile_k + self.num_sfa_tmem_cols = ( + self.cta_tile_shape_mnk[0] // sf_atom_mn + ) * mma_inst_tile_k + self.num_sfb_tmem_cols = ( + self.cta_tile_shape_mnk_sfb[1] // sf_atom_mn + ) * mma_inst_tile_k self.num_sf_tmem_cols = self.num_sfa_tmem_cols + self.num_sfb_tmem_cols - self.num_accumulator_tmem_cols = self.cta_tile_shape_mnk[1] * self.num_acc_stage if not self.overlapping_accum else self.cta_tile_shape_mnk[1] * 2 - self.num_sf_tmem_cols + self.num_accumulator_tmem_cols = ( + self.cta_tile_shape_mnk[1] * self.num_acc_stage + if not self.overlapping_accum + else self.cta_tile_shape_mnk[1] * 2 - self.num_sf_tmem_cols + ) # Only when overlapping_accum is enabled, we need to release accumulator buffer early in epilogue - # Use -1 since at that iteration the pipeline is updated after the tmem -> reg copy - num_subtiles_in_overlap_region = ceil_div(self.num_sf_tmem_cols, self.epi_tile_n) - self.iter_acc_early_release_in_epilogue = num_subtiles_in_overlap_region - 1 + self.iter_acc_early_release_in_epilogue = ( + self.num_sf_tmem_cols // self.epi_tile_n + ) @cute.jit def __call__( self, - a_tensor: cute.Tensor, - b_tensor: cute.Tensor, - sfa_tensor: cute.Tensor, - sfb_tensor: cute.Tensor, - c_tensor: cute.Tensor, + a_ptr: cute.Pointer, + b_ptr: cute.Pointer, + sfa_ptr: cute.Pointer, + sfb_ptr: cute.Pointer, + c_ptr: cute.Pointer, + layouts: cutlass.Constexpr[ + Tuple[tcgen05.OperandMajorMode, tcgen05.OperandMajorMode, utils.LayoutEnum] + ], + problem_mnkl: Tuple[int, int, int, int], max_active_clusters: cutlass.Constexpr, stream: cuda.CUstream, epilogue_op: cutlass.Constexpr = lambda x: x, @@ -423,13 +430,13 @@ class Sm100BlockScaledPersistentDenseGemmKernel: :raises TypeError: If input data types are incompatible with the MMA instruction. """ # Setup static attributes before smem/grid/tma computation - self.a_dtype: Type[cutlass.Numeric] = a_tensor.element_type - self.b_dtype: Type[cutlass.Numeric] = b_tensor.element_type - self.sf_dtype: Type[cutlass.Numeric] = sfa_tensor.element_type - self.c_dtype: Type[cutlass.Numeric] = c_tensor.element_type - self.a_major_mode = utils.LayoutEnum.from_tensor(a_tensor).mma_major_mode() - self.b_major_mode = utils.LayoutEnum.from_tensor(b_tensor).mma_major_mode() - self.c_layout = utils.LayoutEnum.from_tensor(c_tensor) + self.a_dtype: Type[cutlass.Numeric] = a_ptr.value_type + self.b_dtype: Type[cutlass.Numeric] = b_ptr.value_type + self.sf_dtype: Type[cutlass.Numeric] = sfa_ptr.value_type + self.c_dtype: Type[cutlass.Numeric] = c_ptr.value_type + + m, n, k, l = problem_mnkl + self.a_major_mode, self.b_major_mode, self.c_layout = layouts # Check if input data types are compatible with MMA instruction if cutlass.const_expr(self.a_dtype != self.b_dtype): @@ -438,18 +445,37 @@ class Sm100BlockScaledPersistentDenseGemmKernel: # Setup attributes that dependent on gemm inputs self._setup_attributes() + a_layout = cute.make_ordered_layout((m, cute.assume(k, 32), l), order=(0, 1, 2)) + if cutlass.const_expr(self.a_major_mode == tcgen05.OperandMajorMode.K): + a_layout = cute.make_ordered_layout( + (cute.assume(m, 32), k, l), order=(1, 0, 2) + ) + b_layout = cute.make_ordered_layout((n, cute.assume(k, 32), l), order=(0, 1, 2)) + if cutlass.const_expr(self.b_major_mode == tcgen05.OperandMajorMode.K): + b_layout = cute.make_ordered_layout( + (cute.assume(n, 32), k, l), order=(1, 0, 2) + ) + c_layout = cute.make_ordered_layout((cute.assume(m, 32), n, l), order=(0, 1, 2)) + if cutlass.const_expr(self.c_layout == utils.LayoutEnum.ROW_MAJOR): + c_layout = cute.make_ordered_layout( + (m, cute.assume(n, 32), l), order=(1, 0, 2) + ) + a_tensor = cute.make_tensor(a_ptr, a_layout) + b_tensor = cute.make_tensor(b_ptr, b_layout) + c_tensor = cute.make_tensor(c_ptr, c_layout) + # Setup sfa/sfb tensor by filling A/B tensor to scale factor atom layout # ((Atom_M, Rest_M),(Atom_K, Rest_K),RestL) sfa_layout = blockscaled_utils.tile_atom_to_shape_SF( a_tensor.shape, self.sf_vec_size ) - sfa_tensor = cute.make_tensor(sfa_tensor.iterator, sfa_layout) + sfa_tensor = cute.make_tensor(sfa_ptr, sfa_layout) # ((Atom_N, Rest_N),(Atom_K, Rest_K),RestL) sfb_layout = blockscaled_utils.tile_atom_to_shape_SF( b_tensor.shape, self.sf_vec_size ) - sfb_tensor = cute.make_tensor(sfb_tensor.iterator, sfb_layout) + sfb_tensor = cute.make_tensor(sfb_ptr, sfb_layout) tiled_mma = sm100_utils.make_blockscaled_trivial_tiled_mma( self.a_dtype, @@ -539,25 +565,21 @@ class Sm100BlockScaledPersistentDenseGemmKernel: y = cute.ceil_div(tma_tensor_sfb.shape[0][1], 4) new_shape = ( - ( - tma_tensor_sfb.shape[0][0], - ((2, 2), y) - ), + (tma_tensor_sfb.shape[0][0], ((2, 2), y)), tma_tensor_sfb.shape[1], - tma_tensor_sfb.shape[2] + tma_tensor_sfb.shape[2], ) # Use right multiplication for ScaledBasis (3 * x instead of x * 3) x_times_3 = 3 * x new_stride = ( - ( - tma_tensor_sfb.stride[0][0], - ((x, x), x_times_3) - ), + (tma_tensor_sfb.stride[0][0], ((x, x), x_times_3)), tma_tensor_sfb.stride[1], - tma_tensor_sfb.stride[2] + tma_tensor_sfb.stride[2], ) tma_tensor_sfb_new_layout = cute.make_layout(new_shape, stride=new_stride) - tma_tensor_sfb = cute.make_tensor(tma_tensor_sfb.iterator, tma_tensor_sfb_new_layout) + tma_tensor_sfb = cute.make_tensor( + tma_tensor_sfb.iterator, tma_tensor_sfb_new_layout + ) a_copy_size = cute.size_in_bytes(self.a_dtype, a_smem_layout) b_copy_size = cute.size_in_bytes(self.b_dtype, b_smem_layout) @@ -896,7 +918,7 @@ class Sm100BlockScaledPersistentDenseGemmKernel: cute.group_modes(tCgB, 0, 3), ) - # TMA load SFA partition_S/D + # TMA load scaled factor A partition_S/D sfa_cta_layout = a_cta_layout # ((atom_v, rest_v), STAGE) # ((atom_v, rest_v), RestM, RestK, RestL) @@ -910,7 +932,7 @@ class Sm100BlockScaledPersistentDenseGemmKernel: tAsSFA = cute.filter_zeros(tAsSFA) tAgSFA = cute.filter_zeros(tAgSFA) - # TMA load SFB partition_S/D + # TMA load scaled factor B partition_S/D sfb_cta_layout = cute.make_layout( cute.slice_(cluster_layout_sfb_vmnk, (0, None, 0, 0)).shape ) @@ -945,13 +967,13 @@ class Sm100BlockScaledPersistentDenseGemmKernel: tCtAcc_fake.iterator, cute.make_layout( tCtAcc_fake.shape, - stride = ( + stride=( tCtAcc_fake.stride[0], tCtAcc_fake.stride[1], tCtAcc_fake.stride[2], - (256 - self.num_sf_tmem_cols) * tCtAcc_fake.stride[0][1] - ) - ) + (256 - self.num_sf_tmem_cols) * tCtAcc_fake.stride[0][1], + ), + ), ) else: # (MMA, MMA_M, MMA_N, STAGE) @@ -1010,9 +1032,7 @@ class Sm100BlockScaledPersistentDenseGemmKernel: if cutlass.const_expr(self.cta_tile_shape_mnk[1] == 64): slice_n = mma_tile_coord_mnl[1] // 2 # ((atom_v, rest_v), RestK) - tBgSFB_slice = tBgSFB[ - (None, slice_n, None, mma_tile_coord_mnl[2]) - ] + tBgSFB_slice = tBgSFB[(None, slice_n, None, mma_tile_coord_mnl[2])] # Peek (try_wait) AB buffer empty for k_tile = prefetch_k_tile_cnt ab_producer_state.reset_count() @@ -1188,11 +1208,15 @@ class Sm100BlockScaledPersistentDenseGemmKernel: tCtSFB_mma = tCtSFB if cutlass.const_expr(self.cta_tile_shape_mnk[1] == 192): # If this is an ODD tile, shift the TMEM start address for cta_tile_shape_n=192 case by two words (ignores first 64 columns of SFB) - offset = cutlass.Int32(2) if mma_tile_coord_mnl[1] % 2 == 1 else cutlass.Int32(0) + offset = ( + cutlass.Int32(2) + if mma_tile_coord_mnl[1] % 2 == 1 + else cutlass.Int32(0) + ) shifted_ptr = cute.recast_ptr( acc_tmem_ptr - + self.num_accumulator_tmem_cols - + self.num_sfa_tmem_cols + + self.num_accumulator_tmem_cols + + self.num_sfa_tmem_cols + offset, dtype=self.sf_dtype, ) @@ -1201,7 +1225,7 @@ class Sm100BlockScaledPersistentDenseGemmKernel: # Move in increments of 64 columns of SFB offset = cutlass.Int32((mma_tile_coord_mnl[1] % 2) * 2) shifted_ptr = cute.recast_ptr( - acc_tmem_ptr + acc_tmem_ptr + self.num_accumulator_tmem_cols + self.num_sfa_tmem_cols + offset, @@ -1398,7 +1422,7 @@ class Sm100BlockScaledPersistentDenseGemmKernel: # Get accumulator stage index if cutlass.const_expr(self.overlapping_accum): acc_stage_index = acc_consumer_state.phase - reverse_subtile = cutlass.Boolean(True) if acc_stage_index == 0 else cutlass.Boolean(False) + reverse_subtile = True if acc_stage_index == 0 else False else: acc_stage_index = acc_consumer_state.index @@ -1425,7 +1449,11 @@ class Sm100BlockScaledPersistentDenseGemmKernel: real_subtile_idx = subtile_idx if cutlass.const_expr(self.overlapping_accum): if reverse_subtile: - real_subtile_idx = self.cta_tile_shape_mnk[1] // self.epi_tile_n - 1 - subtile_idx + real_subtile_idx = ( + self.cta_tile_shape_mnk[1] // self.epi_tile_n + - 1 + - subtile_idx + ) # # Load accumulator from tensor memory buffer to register # @@ -1459,7 +1487,10 @@ class Sm100BlockScaledPersistentDenseGemmKernel: tRS_sC[(None, None, None, c_buffer)], ) # Fence and barrier to make sure shared memory store is visible to TMA store - cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) self.epilog_sync_barrier.arrive_and_wait() # @@ -1900,9 +1931,9 @@ class Sm100BlockScaledPersistentDenseGemmKernel: def is_valid_layouts( ab_dtype: Type[cutlass.Numeric], c_dtype: Type[cutlass.Numeric], - a_major: str, - b_major: str, - c_major: str, + a_major: Literal["m", "k"], + b_major: Literal["n", "k"], + c_major: Literal["m", "n"], ) -> bool: """ Check if layouts and dtypes are valid combinations @@ -1912,11 +1943,11 @@ class Sm100BlockScaledPersistentDenseGemmKernel: :param c_dtype: The data type of the output tensor :type c_dtype: Type[cutlass.Numeric] :param a_major: The major dimension of the A tensor - :type a_major: str + :type a_major: Literal["m", "k"] :param b_major: The major dimension of the B tensor - :type b_major: str + :type b_major: Literal["n", "k"] :param c_major: The major dimension of the C tensor - :type c_major: str + :type c_major: Literal["m", "n"] :return: True if the layouts are valid, False otherwise :rtype: bool @@ -1976,9 +2007,9 @@ class Sm100BlockScaledPersistentDenseGemmKernel: l: int, ab_dtype: Type[cutlass.Numeric], c_dtype: Type[cutlass.Numeric], - a_major: str, - b_major: str, - c_major: str, + a_major: Literal["m", "k"], + b_major: Literal["n", "k"], + c_major: Literal["m", "n"], ) -> bool: """ Check if the tensor alignment is valid @@ -1996,11 +2027,11 @@ class Sm100BlockScaledPersistentDenseGemmKernel: :param c_dtype: The data type of the output tensor :type c_dtype: Type[cutlass.Numeric] :param a_major: The major axis of the A tensor - :type a_major: str + :type a_major: Literal["m", "k"] :param b_major: The major axis of the B tensor - :type b_major: str + :type b_major: Literal["n", "k"] :param c_major: The major axis of the C tensor - :type c_major: str + :type c_major: Literal["m", "n"] :return: True if the problem shape is valid, False otherwise :rtype: bool @@ -2023,27 +2054,32 @@ class Sm100BlockScaledPersistentDenseGemmKernel: @staticmethod def can_implement( + mnkl: Tuple[int, int, int, int], ab_dtype: Type[cutlass.Numeric], sf_dtype: Type[cutlass.Numeric], - sf_vec_size: int, c_dtype: Type[cutlass.Numeric], + a_major: Literal["m", "k"], + b_major: Literal["n", "k"], + c_major: Literal["m", "n"], + sf_vec_size: int, mma_tiler_mn: Tuple[int, int], cluster_shape_mn: Tuple[int, int], - m: int, - n: int, - k: int, - l: int, - a_major: str, - b_major: str, - c_major: str, ) -> bool: """ Check if the gemm can be implemented + :param mnkl: The problem size as a tuple (M, N, K, L). + :type mnkl: Tuple[int, int, int, int] :param ab_dtype: The data type of the A and B operands :type ab_dtype: Type[cutlass.Numeric] :param sf_dtype: The data type of the scale factor tensor :type sf_dtype: Type[cutlass.Numeric] + :param a_major: The major axis of the A tensor + :type a_major: Literal["m", "k"] + :param b_major: The major axis of the B tensor + :type b_major: Literal["n", "k"] + :param c_major: The major axis of the C tensor + :type c_major: Literal["m", "n"] :param sf_vec_size: The vector size :type sf_vec_size: int :param c_dtype: The data type of the output tensor @@ -2052,24 +2088,11 @@ class Sm100BlockScaledPersistentDenseGemmKernel: :type mma_tiler_mn: Tuple[int, int] :param cluster_shape_mn: The (ClusterM, ClusterN) shape of the CTA cluster :type cluster_shape_mn: Tuple[int, int] - :param m: The number of rows in the A tensor - :type m: int - :param n: The number of columns in the B tensor - :type n: int - :param k: The number of columns in the A tensor - :type k: int - :param l: The number of columns in the C tensor - :type l: int - :param a_major: The major axis of the A tensor - :type a_major: str - :param b_major: The major axis of the B tensor - :type b_major: str - :param c_major: The major axis of the C tensor - :type c_major: str - :return: True if the gemm can be implemented, False otherwise :rtype: bool """ + # Unpack parameters + m, n, k, l = mnkl can_implement = True # Skip unsupported types if not Sm100BlockScaledPersistentDenseGemmKernel.is_valid_dtypes_and_scale_factor_vec_size( @@ -2094,30 +2117,448 @@ class Sm100BlockScaledPersistentDenseGemmKernel: return can_implement +# Helper function to convert scale factor tensor from MKL layout to (32, 4, restM, 4, restK, l) format @cute.jit def cvt_sf_MKL_to_M32x4xrm_K4xrk_L( - sf_ref_tensor: cute.Tensor, - sf_mma_tensor: cute.Tensor, + sf_ref_ptr: cute.Pointer, + sf_mma_ptr: cute.Pointer, + mn: int, + sf_k: int, + l: int, + mma_shape: tuple, ): - """Convert scale factor tensor from MKL layout to mma specification M(32x4xrest_m)xK(4xrest_k)xL layout""" - # sf_mma_tensor has flatten shape (32, 4, rest_m, 4, rest_k, l) - # group to ((32, 4, rest_m), (4, rest_k), l) + mma_permute_order = (3, 4, 1, 5, 2, 0) + permuted_shape = tuple(mma_shape[i] for i in mma_permute_order) + cute_layout = cute.make_ordered_layout(permuted_shape, order=(2, 1, 4, 0, 3, 5)) + + sf_ref_tensor = cute.make_tensor( + sf_ref_ptr, cute.make_layout((mn, sf_k, l), stride=(sf_k, 1, mn * sf_k)) + ) + sf_mma_tensor = cute.make_tensor(sf_mma_ptr, cute_layout) + sf_mma_tensor = cute.group_modes(sf_mma_tensor, 0, 3) sf_mma_tensor = cute.group_modes(sf_mma_tensor, 1, 3) for i in cutlass.range(cute.size(sf_ref_tensor)): mkl_coord = sf_ref_tensor.layout.get_hier_coord(i) sf_mma_tensor[mkl_coord] = sf_ref_tensor[mkl_coord] + pass -def run( +# Helper function for ceil division +def ceil_div(a, b): + return (a + b - 1) // b + + +# Convert scale factor tensors from (m, k, l) to (32, 4, restM, 4, restK, l) format +def create_and_reorder_scale_factor_tensor( + l, mn, k, sf_vec_size, sf_dtype, torch_tensor +): + """ + Create the CUTE-format scale factor tensor on CUDA based on the reference tensor. + """ + sf_k = ceil_div(k, sf_vec_size) + atom_m = (32, 4) + atom_k = 4 + mma_shape = ( + l, # batch size + ceil_div(mn, atom_m[0] * atom_m[1]), + ceil_div(sf_k, atom_k), + atom_m[0], + atom_m[1], + atom_k, + ) + + # Generate a random int8 tensor, then convert to float8_e4m3fn + cute_tensor = torch.ones(mma_shape, dtype=cutlass_torch.dtype(sf_dtype)).permute( + 3, 4, 1, 5, 2, 0 + ) + + # Call the helper function to do layout conversion + cvt_sf_MKL_to_M32x4xrm_K4xrk_L( + make_ptr( + sf_dtype, + torch_tensor.data_ptr(), + cute.AddressSpace.gmem, + assumed_align=32, + ), + make_ptr( + sf_dtype, + cute_tensor.data_ptr(), + cute.AddressSpace.gmem, + assumed_align=32, + ), + mn, + sf_k, + l, + mma_shape, + ) + return cute_tensor.cuda() + + +# Compile the persistent dense blockscaled GEMM operation +def scaled_mm( + gemm_obj: Sm100BlockScaledPersistentDenseGemmKernel, + ab_dtype: Type[cutlass.Numeric], + c_dtype: Type[cutlass.Numeric], + sf_dtype: Type[cutlass.Numeric], + a_major: Literal["m", "k"], + b_major: Literal["n", "k"], + c_major: Literal["m", "n"], + max_active_clusters: cutlass.Constexpr, + stream: cuda.CUstream, + epilogue_op: cutlass.Constexpr = lambda x: x, + options: str = "", +): + # Construct CuTe Pointers + a_ptr = make_ptr(ab_dtype, 0, cute.AddressSpace.gmem, assumed_align=16) + b_ptr = make_ptr(ab_dtype, 0, cute.AddressSpace.gmem, assumed_align=16) + c_ptr = make_ptr(c_dtype, 0, cute.AddressSpace.gmem, assumed_align=16) + sfa_ptr = make_ptr(sf_dtype, 0, cute.AddressSpace.gmem, assumed_align=32) + sfb_ptr = make_ptr(sf_dtype, 0, cute.AddressSpace.gmem, assumed_align=32) + + a_major_mode = ( + tcgen05.OperandMajorMode.K if a_major == "k" else tcgen05.OperandMajorMode.MN + ) + b_major_mode = ( + tcgen05.OperandMajorMode.K if b_major == "k" else tcgen05.OperandMajorMode.MN + ) + c_layout = ( + utils.LayoutEnum.ROW_MAJOR if c_major == "n" else utils.LayoutEnum.COL_MAJOR + ) + return cute.compile( + gemm_obj, + a_ptr, + b_ptr, + sfa_ptr, + sfb_ptr, + c_ptr, + (a_major_mode, b_major_mode, c_layout), + (cutlass.Int32(0), cutlass.Int32(0), cutlass.Int32(0), cutlass.Int32(0)), + max_active_clusters, + stream, + epilogue_op, + options=options, + ) + + +def is_emulated_dtype( + ab_dtype: Type[cutlass.Numeric], + sf_dtype: Type[cutlass.Numeric], + c_dtype: Type[cutlass.Numeric], +) -> bool: + if c_dtype in { + cutlass.Float32, + cutlass.Float16, + cutlass.BFloat16, + }: + if ab_dtype == cutlass.Float4E2M1FN and sf_dtype == cutlass.Float8E4M3FN: + return False + if ab_dtype == cutlass.Float8E4M3FN and sf_dtype == cutlass.Float8E8M0FNU: + return False + + return True + + +# Convert scale factor tensor from MKL layout to blocked layout +def to_blocked(input_matrix): + rows, cols = input_matrix.shape + # Please ensure rows and cols are multiples of 128 and 4 respectively + n_row_blocks = ceil_div(rows, 128) + n_col_blocks = ceil_div(cols, 4) + padded_rows = n_row_blocks * 128 + padded_cols = n_col_blocks * 4 + + # Pad the input matrix if necessary + if padded_rows != rows or padded_cols != cols: + # For FP8 types, convert to float32 for padding, then convert back + original_dtype = input_matrix.dtype + input_float32 = input_matrix.to(torch.float32) + padded = torch.nn.functional.pad( + input_float32, + (0, padded_cols - cols, 0, padded_rows - rows), + mode="constant", + value=0, + ) + # Convert back to original dtype if needed + if original_dtype != input_float32.dtype: + padded = padded.to(original_dtype) + else: + padded = input_matrix + blocks = padded.view(n_row_blocks, 128, n_col_blocks, 4).permute(0, 2, 1, 3) + rearranged = blocks.reshape(-1, 4, 32, 4).transpose(1, 2).reshape(-1, 32, 16) + return rearranged.flatten() + + +# Reference implementation of the persistent dense blockscaled GEMM operation (emulated version) +def reference_scaled_mm_emulated( + a: torch.Tensor, + b: torch.Tensor, + sfa: torch.Tensor, + sfb: torch.Tensor, + c: torch.Tensor, + mnkl: Tuple[int, int, int, int], + sf_vec_size: int, + c_dtype: Type[cutlass.Numeric], +): + m, n, k, l = mnkl + sfa_expanded = ( + torch.repeat_interleave(sfa, sf_vec_size, dim=1)[:, :k, :] + .to(dtype=torch.float32) + .cuda() + ) + sfb_expanded = ( + torch.repeat_interleave(sfb, sf_vec_size, dim=1)[:, :k, :] + .to(dtype=torch.float32) + .cuda() + ) + res_a = torch.einsum("mkl,mkl->mkl", a, sfa_expanded) + res_b = torch.einsum("nkl,nkl->nkl", b, sfb_expanded) + # Cast res_a and res_b to float32 for einsum to avoid NotImplementedError on 'Byte' + ref = torch.einsum("mkl,nkl->mnl", res_a, res_b) + c_ref = ref.to(dtype=cutlass_torch.dtype(c_dtype)) + return c_ref + + +# Reference implementation of the persistent dense blockscaled GEMM operation (non-emulated version) +def reference_scaled_mm( + a: torch.Tensor, + b: torch.Tensor, + sfa: torch.Tensor, + sfb: torch.Tensor, + c: torch.Tensor, + mnkl: Tuple[int, int, int, int], + c_dtype: Type[cutlass.Numeric], +): + m, n, k, l = mnkl + c_ref = torch.clone(c) + for l_idx in range(l): + # Convert the scale factor tensor to blocked format + scale_a = to_blocked(sfa[:, :, l_idx]) + scale_b = to_blocked(sfb[:, :, l_idx]) + # Ensure a_slice is row-major (M, K) with stride (K, 1) + a_slice = a[:, :, l_idx].contiguous() + # Ensure b_slice is row-major (N, K) so that transpose gives column-major (K, N) + b_slice = b[:, :, l_idx].contiguous() + # (m, k) @ (n, k).T -> (m, n) + res = torch._scaled_mm( + a_slice, + b_slice.transpose(0, 1), + scale_a.cuda(), + scale_b.cuda(), + bias=None, + out_dtype=c_ref.dtype, + ) + c_ref[:, :, l_idx] = res + return c_ref + + +# Construct CuTe Pointers for the persistent dense blockscaled GEMM operation (emulated version) +def construct_cute_pointers_emulated( + a: torch.Tensor, + b: torch.Tensor, + sfa: torch.Tensor, + sfb: torch.Tensor, + c: torch.Tensor, + ab_dtype: Type[cutlass.Numeric], + sf_dtype: Type[cutlass.Numeric], + c_dtype: Type[cutlass.Numeric], +): + a_cute, _ = cutlass_torch.cute_tensor_like( + a.cpu(), + ab_dtype, + is_dynamic_layout=True, + assumed_align=16, + ) + a_cute = cutlass_torch.convert_cute_tensor( + a, + a_cute, + ab_dtype, + is_dynamic_layout=True, + ) + b_cute, _ = cutlass_torch.cute_tensor_like( + b.cpu(), + ab_dtype, + is_dynamic_layout=True, + assumed_align=16, + ) + b_cute = cutlass_torch.convert_cute_tensor( + b, + b_cute, + ab_dtype, + is_dynamic_layout=True, + ) + a_ptr = a_cute.iterator + b_ptr = b_cute.iterator + + sfa_ptr = make_ptr( + sf_dtype, sfa.data_ptr(), cute.AddressSpace.gmem, assumed_align=32 + ) + sfb_ptr = make_ptr( + sf_dtype, sfb.data_ptr(), cute.AddressSpace.gmem, assumed_align=32 + ) + c_ptr = make_ptr(c_dtype, c.data_ptr(), cute.AddressSpace.gmem, assumed_align=16) + return a_ptr, b_ptr, c_ptr, sfa_ptr, sfb_ptr, a_cute, b_cute + + +# Construct CuTe Pointers for the persistent dense blockscaled GEMM operation (non-emulated version) +def construct_cute_pointers( + a: torch.Tensor, + b: torch.Tensor, + sfa: torch.Tensor, + sfb: torch.Tensor, + c: torch.Tensor, + ab_dtype: Type[cutlass.Numeric], + sf_dtype: Type[cutlass.Numeric], + c_dtype: Type[cutlass.Numeric], +): + a_ptr = make_ptr(ab_dtype, a.data_ptr(), cute.AddressSpace.gmem, assumed_align=16) + b_ptr = make_ptr(ab_dtype, b.data_ptr(), cute.AddressSpace.gmem, assumed_align=16) + sfa_ptr = make_ptr( + sf_dtype, sfa.data_ptr(), cute.AddressSpace.gmem, assumed_align=32 + ) + sfb_ptr = make_ptr( + sf_dtype, sfb.data_ptr(), cute.AddressSpace.gmem, assumed_align=32 + ) + c_ptr = make_ptr(c_dtype, c.data_ptr(), cute.AddressSpace.gmem, assumed_align=16) + return a_ptr, b_ptr, c_ptr, sfa_ptr, sfb_ptr + + +# Use uint8 and uint32 to emulate unsupported +# dtype in torch +def prepare_tensors_emulated( mnkl: Tuple[int, int, int, int], ab_dtype: Type[cutlass.Numeric], sf_dtype: Type[cutlass.Numeric], sf_vec_size: int, c_dtype: Type[cutlass.Numeric], - a_major: str, - b_major: str, - c_major: str, + a_major: Literal["m", "k"], + b_major: Literal["n", "k"], + c_major: Literal["m", "n"], +): + m, n, k, l = mnkl + sf_k = ceil_div(k, sf_vec_size) + + # Create tensor SFA/SFB with values in [1, 3) + sfa = ( + torch.randint(0, 3, (l, m, sf_k), dtype=torch.uint8) + .permute(1, 2, 0) + .to(dtype=cutlass_torch.dtype(sf_dtype)) + ) + sfb = ( + torch.randint(0, 3, (l, n, sf_k), dtype=torch.uint8) + .permute(1, 2, 0) + .to(dtype=cutlass_torch.dtype(sf_dtype)) + ) + + # Create tensor A/B with values in [0, 2) + if a_major == "k": + a = torch.randint(-2, 2, (l, m, k), dtype=torch.float32, device="cuda").permute( + 1, 2, 0 + ) + else: + a = torch.randint(-2, 2, (l, k, m), dtype=torch.float32, device="cuda").permute( + 2, 1, 0 + ) + if b_major == "k": + b = torch.randint(-2, 2, (l, n, k), dtype=torch.float32, device="cuda").permute( + 1, 2, 0 + ) + else: + b = torch.randint(-2, 2, (l, k, n), dtype=torch.float32, device="cuda").permute( + 2, 1, 0 + ) + if c_major == "n": + c = torch.empty( + (l, m, n), dtype=cutlass_torch.dtype(c_dtype), device="cuda" + ).permute(1, 2, 0) + else: + c = torch.empty( + (l, n, m), dtype=cutlass_torch.dtype(c_dtype), device="cuda" + ).permute(2, 1, 0) + return a, b, c, sfa, sfb + + +def prepare_tensors( + mnkl: Tuple[int, int, int, int], + ab_dtype: Type[cutlass.Numeric], + sf_dtype: Type[cutlass.Numeric], + sf_vec_size: int, + c_dtype: Type[cutlass.Numeric], + a_major: Literal["m", "k"], + b_major: Literal["n", "k"], + c_major: Literal["m", "n"], +): + m, n, k, l = mnkl + + if ab_dtype == cutlass.Float4E2M1FN: + # Using int8 for torch.float4_e2m1fn_x2 tensor allocation + # Thus the size of k needs to be halved in this case. + k_fct = 2 + else: + k_fct = 1 + + sf_k = ceil_div(k, sf_vec_size) + + # Create tensor SFA/SFB + sfa = ( + torch.randint(0, 3, (l, m, sf_k), dtype=torch.uint8) + .permute(1, 2, 0) + .to(dtype=cutlass_torch.dtype(sf_dtype)) + ) + sfb = ( + torch.randint(0, 3, (l, n, sf_k), dtype=torch.uint8) + .permute(1, 2, 0) + .to(dtype=cutlass_torch.dtype(sf_dtype)) + ) + + # Create tensor A/B/C + if a_major == "k": + a = torch.randint( + -2, 2, (l, m, k // k_fct), dtype=torch.int8, device="cuda" + ).permute(1, 2, 0) + else: + a = torch.randint(-2, 2, (l, k, m), dtype=torch.int8, device="cuda").permute( + 2, 1, 0 + ) + if b_major == "k": + b = torch.randint( + -2, 2, (l, n, k // k_fct), dtype=torch.int8, device="cuda" + ).permute(1, 2, 0) + else: + b = torch.randint(-2, 2, (l, k, n), dtype=torch.int8, device="cuda").permute( + 2, 1, 0 + ) + if c_major == "n": + c = torch.randint( + -2, 2, (l, m, n), dtype=cutlass_torch.dtype(c_dtype), device="cuda" + ).permute(1, 2, 0) + else: + c = torch.randint( + -2, 2, (l, n, m), dtype=cutlass_torch.dtype(c_dtype), device="cuda" + ).permute(2, 1, 0) + + if ab_dtype == cutlass.Float4E2M1FN: + a = a.view(dtype=torch.float4_e2m1fn_x2) + b = b.view(dtype=torch.float4_e2m1fn_x2) + else: + a = a.to(dtype=cutlass_torch.dtype(ab_dtype)) + b = b.to(dtype=cutlass_torch.dtype(ab_dtype)) + + c = c.to(dtype=cutlass_torch.dtype(c_dtype)) + return a, b, c, sfa, sfb + + +# This will show how to covert torch tensor +# and pass to CuTe kernel +def run_scaled_mm( + mnkl: Tuple[int, int, int, int], + ab_dtype: Type[cutlass.Numeric], + sf_dtype: Type[cutlass.Numeric], + sf_vec_size: int, + c_dtype: Type[cutlass.Numeric], + a_major: Literal["m", "k"], + b_major: Literal["n", "k"], + c_major: Literal["m", "n"], mma_tiler_mn: Tuple[int, int], cluster_shape_mn: Tuple[int, int], tolerance: float = 1e-01, @@ -2127,7 +2568,7 @@ def run( use_cold_l2: bool = False, **kwargs, ): - """Execute a persistent batched dense blockscaled GEMM operation on Blackwell architecture with performance benchmarking. + """Execute a persistent batched dense blockscaled GEMM operation on Blackwell architecture with performance benchmarking (non-emulated dtypes). This function prepares input tensors, configures and launches the persistent GEMM kernel, optionally performs reference validation, and benchmarks the execution performance. @@ -2143,7 +2584,7 @@ def run( :param c_dtype: Data type for output tensor C :type c_dtype: Type[cutlass.Numeric] :param a_major/b_major/c_major: Memory layout of tensor A/B/C - :type a_major/b_major/c_major: str + :type a_major/b_major/c_major: Literal["m", "k", "n"] :param mma_tiler_mn: MMA tiling size. :type mma_tiler_mn: Tuple[int, int] :param cluster_shape_mn: Cluster shape. @@ -2178,21 +2619,25 @@ def run( # Unpack parameters m, n, k, l = mnkl - # Skip unsupported testcase - if not Sm100BlockScaledPersistentDenseGemmKernel.can_implement( - ab_dtype, - sf_dtype, + # Configure gemm kernel + gemm = Sm100BlockScaledPersistentDenseGemmKernel( sf_vec_size, - c_dtype, mma_tiler_mn, cluster_shape_mn, - m, - n, - k, - l, + ) + + # Skip unsupported testcase + if not gemm.can_implement( + mnkl, + ab_dtype, + sf_dtype, + c_dtype, a_major, b_major, c_major, + sf_vec_size, + mma_tiler_mn, + cluster_shape_mn, ): raise TypeError( f"Unsupported testcase {ab_dtype}, {sf_dtype}, {sf_vec_size}, {c_dtype}, {mma_tiler_mn}, {cluster_shape_mn}, {m}, {n}, {k}, {l}, {a_major}, {b_major}, {c_major}" @@ -2203,122 +2648,197 @@ def run( torch.manual_seed(1111) - # Create tensor A/B/C - a_ref = cutlass_torch.matrix(l, m, k, a_major == "m", cutlass.Float32) - b_ref = cutlass_torch.matrix(l, n, k, b_major == "n", cutlass.Float32) - c_ref = cutlass_torch.matrix(l, m, n, c_major == "m", cutlass.Float32) + # Get current CUDA stream from PyTorch + torch_stream = torch.cuda.current_stream() + # Get the raw stream pointer as a CUstream + current_stream = cuda.CUstream(torch_stream.cuda_stream) - a_tensor, a_torch = cutlass_torch.cute_tensor_like( - a_ref, ab_dtype, is_dynamic_layout=True, assumed_align=16 - ) - b_tensor, b_torch = cutlass_torch.cute_tensor_like( - b_ref, ab_dtype, is_dynamic_layout=True, assumed_align=16 - ) - c_tensor, c_torch = cutlass_torch.cute_tensor_like( - c_ref, c_dtype, is_dynamic_layout=True, assumed_align=16 + # Check if configuration can be implemented + max_active_clusters = utils.HardwareInfo().get_max_active_clusters( + cluster_shape_mn[0] * cluster_shape_mn[1] ) - # Mark tensor with element divisibility for 16B alignment - a_tensor.mark_compact_shape_dynamic( - mode=1 if a_major == "k" else 0, - stride_order=(2, 0, 1) if a_major == "k" else (2, 1, 0), - divisibility=32 if ab_dtype == cutlass.Float4E2M1FN else 16, - ) - b_tensor.mark_compact_shape_dynamic( - mode=1 if b_major == "k" else 0, - stride_order=(2, 0, 1) if b_major == "k" else (2, 1, 0), - divisibility=32 if ab_dtype == cutlass.Float4E2M1FN else 16, - ) - c_tensor.mark_compact_shape_dynamic( - mode=1 if c_major == "n" else 0, - stride_order=(2, 0, 1) if c_major == "n" else (2, 1, 0), - divisibility=32 if ab_dtype == cutlass.Float4E2M1FN else 16, + # Compile gemm kernel with fake tensors + compiled_gemm = scaled_mm( + gemm, + ab_dtype, + c_dtype, + sf_dtype, + a_major, + b_major, + c_major, + max_active_clusters, + current_stream, + options=f"--opt-level 2", ) - # Create scale factor tensor SFA/SFB - def create_scale_factor_tensor(l, mn, k, sf_vec_size, dtype): - sf_k = ceil_div(k, sf_vec_size) - ref_shape = (l, mn, sf_k) + # Create Torch Tensors for A, scale factor A, B, scale factor B, C + a, b, c, sfa, sfb = prepare_tensors( + mnkl, ab_dtype, sf_dtype, sf_vec_size, c_dtype, a_major, b_major, c_major + ) + # Reorder scale factor tensors to (32, 4, restM, 4, restK, l) format + sfa_reordered = create_and_reorder_scale_factor_tensor( + l, m, k, sf_vec_size, sf_dtype, sfa + ) + sfb_reordered = create_and_reorder_scale_factor_tensor( + l, n, k, sf_vec_size, sf_dtype, sfb + ) + # Construct CuTe Pointers + a_ptr, b_ptr, c_ptr, sfa_ptr, sfb_ptr = construct_cute_pointers( + a, + b, + sfa_reordered, + sfb_reordered, + c, + ab_dtype, + sf_dtype, + c_dtype, + ) - atom_m = (32, 4) - atom_k = 4 - mma_shape = ( - l, - ceil_div(mn, atom_m[0] * atom_m[1]), - ceil_div(sf_k, atom_k), - atom_m[0], - atom_m[1], - atom_k, + # Compute reference result + if not skip_ref_check: + # Execute kernel once for reference checking + compiled_gemm( + a_ptr, b_ptr, sfa_ptr, sfb_ptr, c_ptr, (m, n, k, l), current_stream + ) + c_ref = reference_scaled_mm(a, b, sfa, sfb, c, (m, n, k, l), c_dtype) + if c_dtype in (cutlass.Float8E5M2, cutlass.Float8E4M3FN): + # Rtol=0.001 and atol=0.1 are not supported for bitwise comparison of + # low dimensional floats. Please use rtol=0.0 and atol=0.0. + tolerance = 0.0 + torch.testing.assert_close(c, c_ref, atol=tolerance, rtol=tolerance) + + def generate_inputs(): + a, b, c, sfa, sfb = prepare_tensors( + mnkl, + ab_dtype, + sf_dtype, + sf_vec_size, + c_dtype, + a_major, + b_major, + c_major, + ) + # Reorder scale factor tensors to (32, 4, restM, 4, restK, l) format + sfa_reordered = create_and_reorder_scale_factor_tensor( + l, m, k, sf_vec_size, sf_dtype, sfa + ) + sfb_reordered = create_and_reorder_scale_factor_tensor( + l, n, k, sf_vec_size, sf_dtype, sfb + ) + # Construct CuTe Pointers + a_ptr, b_ptr, c_ptr, sfa_ptr, sfb_ptr = construct_cute_pointers( + a, + b, + sfa_reordered, + sfb_reordered, + c, + ab_dtype, + sf_dtype, + c_dtype, + ) + jit_args = cute.testing.JitArguments( + a_ptr, b_ptr, sfa_ptr, sfb_ptr, c_ptr, (m, n, k, l), current_stream + ) + # Keep references to external variables (e.g., Torch tensors when taking a view) + jit_args.add_to_scope([a, b, sfa_reordered, sfb_reordered, c]) + return jit_args + + workspace_count = 1 + if use_cold_l2: + one_workspace_bytes = ( + a.numel() * a.element_size() + + b.numel() * b.element_size() + + sfa.numel() * sfa.element_size() + + sfb.numel() * sfb.element_size() + + c.numel() * c.element_size() + ) + workspace_count = cute.testing.get_workspace_count( + one_workspace_bytes, warmup_iterations, iterations ) - ref_permute_order = (1, 2, 0) - mma_permute_order = (3, 4, 1, 5, 2, 0) - - # Create f32 ref torch tensor (cpu) - ref_f32_torch_tensor_cpu = cutlass_torch.create_and_permute_torch_tensor( - ref_shape, - torch.float32, - permute_order=ref_permute_order, - init_type=cutlass_torch.TensorInitType.RANDOM, - init_config=cutlass_torch.RandomInitConfig( - min_val=1, - max_val=3, - ), - ) - - # Create f32 cute torch tensor (cpu) - cute_f32_torch_tensor_cpu = cutlass_torch.create_and_permute_torch_tensor( - mma_shape, - torch.float32, - permute_order=mma_permute_order, - init_type=cutlass_torch.TensorInitType.RANDOM, - init_config=cutlass_torch.RandomInitConfig( - min_val=0, - max_val=1, - ), - ) - - # convert ref f32 tensor to cute f32 tensor - cvt_sf_MKL_to_M32x4xrm_K4xrk_L( - from_dlpack(ref_f32_torch_tensor_cpu), - from_dlpack(cute_f32_torch_tensor_cpu), - ) - cute_f32_torch_tensor = cute_f32_torch_tensor_cpu.cuda() - - # reshape makes memory contiguous - ref_f32_torch_tensor_cpu = ( - ref_f32_torch_tensor_cpu.permute(2, 0, 1) - .unsqueeze(-1) - .expand(l, mn, sf_k, sf_vec_size) - .reshape(l, mn, sf_k * sf_vec_size) - .permute(*ref_permute_order) - ) - # prune to mkl for reference check. - ref_f32_torch_tensor_cpu = ref_f32_torch_tensor_cpu[:, :k, :] - - # Create dtype cute torch tensor (cpu) - cute_tensor, cute_torch_tensor = cutlass_torch.cute_tensor_like( - cute_f32_torch_tensor_cpu, - dtype, - is_dynamic_layout=True, - assumed_align=16, - ) - - # Convert f32 cute tensor to dtype cute tensor - cute_tensor = cutlass_torch.convert_cute_tensor( - cute_f32_torch_tensor, - cute_tensor, - dtype, - is_dynamic_layout=True, - ) - return ref_f32_torch_tensor_cpu, cute_tensor, cute_torch_tensor - - sfa_ref, sfa_tensor, sfa_torch = create_scale_factor_tensor( - l, m, k, sf_vec_size, sf_dtype - ) - sfb_ref, sfb_tensor, sfb_torch = create_scale_factor_tensor( - l, n, k, sf_vec_size, sf_dtype + exec_time = cute.testing.benchmark( + compiled_gemm, + workspace_generator=generate_inputs, + workspace_count=workspace_count, + stream=current_stream, + warmup_iterations=warmup_iterations, + iterations=iterations, ) + return exec_time # Return execution time in microseconds + + +# This is to compatible with the other narrow +# precision combinations are not supported in either +# torch or dlpack. For example, Float4E2M1FN with Float8E8M0FNU. +def run_scaled_mm_with_emulated_dtype( + mnkl: Tuple[int, int, int, int], + ab_dtype: Type[cutlass.Numeric], + sf_dtype: Type[cutlass.Numeric], + sf_vec_size: int, + c_dtype: Type[cutlass.Numeric], + a_major: Literal["m", "k"], + b_major: Literal["n", "k"], + c_major: Literal["m", "n"], + mma_tiler_mn: Tuple[int, int], + cluster_shape_mn: Tuple[int, int], + tolerance: float = 1e-01, + warmup_iterations: int = 0, + iterations: int = 1, + skip_ref_check: bool = False, + use_cold_l2: bool = False, + **kwargs, +): + """Execute a persistent batched dense blockscaled GEMM operation on Blackwell architecture with performance benchmarking (emulated dtypes). + + This function prepares input tensors, configures and launches the persistent GEMM kernel, + optionally performs reference validation, and benchmarks the execution performance. + + :param mnkl: Problem size (M, N, K, L) + :type mnkl: Tuple[int, int, int, int] + :param ab_dtype: Data type for input tensors A and B + :type ab_dtype: Type[cutlass.Numeric] + :param sf_dtype: Data type for scale factor tensor + :type sf_dtype: Type[cutlass.Numeric] + :param sf_vec_size: Vector size for scale factor tensor + :type sf_vec_size: int + :param c_dtype: Data type for output tensor C + :type c_dtype: Type[cutlass.Numeric] + :param a_major/b_major/c_major: Memory layout of tensor A/B/C + :type a_major/b_major/c_major: Literal["m", "n","k"] + :param mma_tiler_mn: MMA tiling size. + :type mma_tiler_mn: Tuple[int, int] + :param cluster_shape_mn: Cluster shape. + :type cluster_shape_mn: Tuple[int, int] + :param tolerance: Tolerance value for reference validation comparison, defaults to 1e-01 + :type tolerance: float, optional + :param warmup_iterations: Number of warmup iterations before benchmarking, defaults to 0 + :type warmup_iterations: int, optional + :param iterations: Number of benchmark iterations to run, defaults to 1 + :type iterations: int, optional + :param skip_ref_check: Whether to skip reference result validation, defaults to False + :type skip_ref_check: bool, optional + :param use_cold_l2: Whether to use circular buffer strategy to ensure cold L2 cache, defaults to False + :type use_cold_l2: bool, optional + :raises RuntimeError: If CUDA GPU is not available + :raises ValueError: If the configuration is invalid or unsupported by the kernel + :return: Execution time of the GEMM kernel + :rtype: float + """ + print("Running Sm100 Persistent Dense BlockScaled GEMM test (Emulated) with:") + print(f"mnkl: {mnkl}") + print(f"AB dtype: {ab_dtype}, SF dtype: {sf_dtype}, SF Vec size: {sf_vec_size}") + print(f"C dtype: {c_dtype}") + print(f"Matrix majors - A: {a_major}, B: {b_major}, C: {c_major}") + print(f"Mma Tiler (M, N): {mma_tiler_mn}, Cluster Shape (M, N): {cluster_shape_mn}") + print(f"Tolerance: {tolerance}") + print(f"Warmup iterations: {warmup_iterations}") + print(f"Iterations: {iterations}") + print(f"Skip reference checking: {skip_ref_check}") + print(f"Use cold L2: {'True' if use_cold_l2 else 'False'}") + + # Unpack parameters + m, n, k, l = mnkl # Configure gemm kernel gemm = Sm100BlockScaledPersistentDenseGemmKernel( @@ -2327,110 +2847,139 @@ def run( cluster_shape_mn, ) - # Compute max active clusters on current device - hardware_info = cutlass.utils.HardwareInfo() - max_active_clusters = hardware_info.get_max_active_clusters( + # Skip unsupported testcase + if not gemm.can_implement( + mnkl, + ab_dtype, + sf_dtype, + c_dtype, + a_major, + b_major, + c_major, + sf_vec_size, + mma_tiler_mn, + cluster_shape_mn, + ): + raise TypeError( + f"Unsupported testcase {ab_dtype}, {sf_dtype}, {sf_vec_size}, {c_dtype}, {mma_tiler_mn}, {cluster_shape_mn}, {m}, {n}, {k}, {l}, {a_major}, {b_major}, {c_major}" + ) + + if not torch.cuda.is_available(): + raise RuntimeError("GPU is required to run this example!") + + torch.manual_seed(1111) + + # Get current CUDA stream from PyTorch + torch_stream = torch.cuda.current_stream() + # Get the raw stream pointer as a CUstream + current_stream = cuda.CUstream(torch_stream.cuda_stream) + + # Check if configuration can be implemented + max_active_clusters = utils.HardwareInfo().get_max_active_clusters( cluster_shape_mn[0] * cluster_shape_mn[1] ) - # Initialize Stream - current_stream = cutlass_torch.default_stream() - - # Compile gemm kernel - compiled_gemm = cute.compile( + # Compile gemm kernel with fake tensors + compiled_gemm = scaled_mm( gemm, - a_tensor, - b_tensor, - sfa_tensor, - sfb_tensor, - c_tensor, + ab_dtype, + c_dtype, + sf_dtype, + a_major, + b_major, + c_major, max_active_clusters, current_stream, options=f"--opt-level 2", ) + # Create Torch Tensors for A, scale factor A, B, scale factor B, C + a, b, c, sfa, sfb = prepare_tensors_emulated( + mnkl, ab_dtype, sf_dtype, sf_vec_size, c_dtype, a_major, b_major, c_major + ) + # Reorder scale factor tensors to (32, 4, restM, 4, restK, l) format + sfa_reordered = create_and_reorder_scale_factor_tensor( + l, m, k, sf_vec_size, sf_dtype, sfa + ) + sfb_reordered = create_and_reorder_scale_factor_tensor( + l, n, k, sf_vec_size, sf_dtype, sfb + ) + # Construct CuTe Pointers + a_ptr, b_ptr, c_ptr, sfa_ptr, sfb_ptr, a_cute, b_cute = ( + construct_cute_pointers_emulated( + a, + b, + sfa_reordered, + sfb_reordered, + c, + ab_dtype, + sf_dtype, + c_dtype, + ) + ) + # Compute reference result if not skip_ref_check: # Execute kernel once for reference checking compiled_gemm( - a_tensor, b_tensor, sfa_tensor, sfb_tensor, c_tensor, current_stream + a_ptr, b_ptr, sfa_ptr, sfb_ptr, c_ptr, (m, n, k, l), current_stream ) - print("Verifying results...") - res_a = torch.einsum("mkl,mkl->mkl", a_ref, sfa_ref) - res_b = torch.einsum("nkl,nkl->nkl", b_ref, sfb_ref) - ref = torch.einsum("mkl,nkl->mnl", res_a, res_b) - - # Convert c back to f32 for comparison. - c_ref_device = c_ref.cuda() - cute.testing.convert( - c_tensor, - from_dlpack(c_ref_device, assumed_align=16).mark_layout_dynamic( - leading_dim=(1 if c_major == "n" else 0) - ), + c_ref = reference_scaled_mm_emulated( + a, b, sfa, sfb, c, (m, n, k, l), sf_vec_size, c_dtype ) - c_ref = c_ref_device.cpu() + if c_dtype in (cutlass.Float8E5M2, cutlass.Float8E4M3FN): + # Rtol=0.001 and atol=0.1 are not supported for bitwise comparison of + # low dimensional floats. Please use rtol=0.0 and atol=0.0. + tolerance = 0.0 + torch.testing.assert_close(c, c_ref, atol=tolerance, rtol=tolerance) - if c_dtype in (cutlass.Float32, cutlass.Float16, cutlass.BFloat16): - torch.testing.assert_close(c_ref, ref, atol=tolerance, rtol=1e-02) - elif c_dtype in (cutlass.Float8E5M2, cutlass.Float8E4M3FN): - # Convert ref : f32 -> f8 -> f32 - ref_f8_ = torch.empty(*(l, m, n), dtype=torch.uint8, device="cuda").permute( - 1, 2, 0 + def generate_inputs(): + a, b, c, sfa, sfb = prepare_tensors_emulated( + mnkl, + ab_dtype, + sf_dtype, + sf_vec_size, + c_dtype, + a_major, + b_major, + c_major, + ) + # Reorder scale factor tensors to (32, 4, restM, 4, restK, l) format + sfa_reordered = create_and_reorder_scale_factor_tensor( + l, m, k, sf_vec_size, sf_dtype, sfa + ) + sfb_reordered = create_and_reorder_scale_factor_tensor( + l, n, k, sf_vec_size, sf_dtype, sfb + ) + # Construct CuTe Pointers + a_ptr, b_ptr, c_ptr, sfa_ptr, sfb_ptr, a_cute, b_cute = ( + construct_cute_pointers_emulated( + a, + b, + sfa_reordered, + sfb_reordered, + c, + ab_dtype, + sf_dtype, + c_dtype, ) - ref_f8 = from_dlpack(ref_f8_, assumed_align=16).mark_layout_dynamic( - leading_dim=1 - ) - ref_f8.element_type = c_dtype - ref_device = ref.permute(2, 0, 1).contiguous().permute(1, 2, 0).cuda() - ref_tensor = from_dlpack(ref_device, assumed_align=16).mark_layout_dynamic( - leading_dim=1 - ) - cute.testing.convert(ref_tensor, ref_f8) - cute.testing.convert(ref_f8, ref_tensor) - ref = ref_device.cpu() - torch.testing.assert_close(c_ref, ref, atol=tolerance, rtol=1e-02) - def generate_tensors(): - a_tensor, _ = cutlass_torch.cute_tensor_like( - a_ref, ab_dtype, is_dynamic_layout=True, assumed_align=16 ) - b_tensor, _ = cutlass_torch.cute_tensor_like( - b_ref, ab_dtype, is_dynamic_layout=True, assumed_align=16 - ) - c_tensor, _ = cutlass_torch.cute_tensor_like( - c_ref, c_dtype, is_dynamic_layout=True, assumed_align=16 + jit_args = cute.testing.JitArguments( + a_ptr, b_ptr, sfa_ptr, sfb_ptr, c_ptr, (m, n, k, l), current_stream ) + # Keep references to external variables (e.g., Torch tensors when taking a view) + jit_args.add_to_scope([a, b, sfa_reordered, sfb_reordered, c, a_cute, b_cute]) + return jit_args - # Mark tensor to be byte aligned - a_tensor.mark_compact_shape_dynamic( - mode=1 if a_major == "k" else 0, - stride_order=(2, 0, 1) if a_major == "k" else (2, 1, 0), - divisibility=2 if ab_dtype == cutlass.Float4E2M1FN else 1, - ) - b_tensor.mark_compact_shape_dynamic( - mode=1 if b_major == "k" else 0, - stride_order=(2, 0, 1) if b_major == "k" else (2, 1, 0), - divisibility=2 if ab_dtype == cutlass.Float4E2M1FN else 1, - ) - c_tensor.mark_compact_shape_dynamic( - mode=1 if c_major == "n" else 0, - stride_order=(2, 0, 1) if c_major == "n" else (2, 1, 0), - divisibility=2 if c_dtype == cutlass.Float4E2M1FN else 1, - ) - - _, sfa_tensor, _ = create_scale_factor_tensor(l, m, k, sf_vec_size, sf_dtype) - _, sfb_tensor, _ = create_scale_factor_tensor(l, n, k, sf_vec_size, sf_dtype) - return cute.testing.JitArguments( - a_tensor, b_tensor, sfa_tensor, sfb_tensor, c_tensor, current_stream - ) workspace_count = 1 if use_cold_l2: one_workspace_bytes = ( - a_torch.numel() * a_torch.element_size() - + b_torch.numel() * b_torch.element_size() - + sfa_torch.numel() * sfa_torch.element_size() - + sfb_torch.numel() * sfb_torch.element_size() - + c_torch.numel() * c_torch.element_size() + a.numel() * a.element_size() + + b.numel() * b.element_size() + + sfa.numel() * sfa.element_size() + + sfb.numel() * sfb.element_size() + + c.numel() * c.element_size() ) workspace_count = cute.testing.get_workspace_count( one_workspace_bytes, warmup_iterations, iterations @@ -2438,16 +2987,78 @@ def run( exec_time = cute.testing.benchmark( compiled_gemm, - workspace_generator=generate_tensors, + workspace_generator=generate_inputs, workspace_count=workspace_count, stream=current_stream, warmup_iterations=warmup_iterations, iterations=iterations, ) - return exec_time # Return execution time in microseconds +def run( + mnkl: Tuple[int, int, int, int], + ab_dtype: Type[cutlass.Numeric], + sf_dtype: Type[cutlass.Numeric], + sf_vec_size: int, + c_dtype: Type[cutlass.Numeric], + a_major: Literal["m", "k"], + b_major: Literal["n", "k"], + c_major: Literal["m", "n"], + mma_tiler_mn: Tuple[int, int], + cluster_shape_mn: Tuple[int, int], + tolerance: float = 1e-01, + warmup_iterations: int = 0, + iterations: int = 1, + skip_ref_check: bool = False, + use_cold_l2: bool = False, + **kwargs, +): + """ + Execute the appropriate GEMM function based on dtype. + + Routes to either run_scaled_mm_with_emulated_dtype or run_scaled_mm + depending on whether the dtypes require emulation. + """ + if is_emulated_dtype(ab_dtype, sf_dtype, c_dtype): + exec_time = run_scaled_mm_with_emulated_dtype( + mnkl, + ab_dtype, + sf_dtype, + sf_vec_size, + c_dtype, + a_major, + b_major, + c_major, + mma_tiler_mn, + cluster_shape_mn, + tolerance, + warmup_iterations, + iterations, + skip_ref_check, + use_cold_l2, + ) + else: + exec_time = run_scaled_mm( + mnkl, + ab_dtype, + sf_dtype, + sf_vec_size, + c_dtype, + a_major, + b_major, + c_major, + mma_tiler_mn, + cluster_shape_mn, + tolerance, + warmup_iterations, + iterations, + skip_ref_check, + use_cold_l2, + ) + return exec_time + + if __name__ == "__main__": def parse_comma_separated_ints(s: str) -> Tuple[int, ...]: @@ -2481,7 +3092,7 @@ if __name__ == "__main__": help="Cluster shape (comma-separated)", ) parser.add_argument("--ab_dtype", type=cutlass.dtype, default=cutlass.Float4E2M1FN) - parser.add_argument("--sf_dtype", type=cutlass.dtype, default=cutlass.Float8E8M0FNU) + parser.add_argument("--sf_dtype", type=cutlass.dtype, default=cutlass.Float8E4M3FN) parser.add_argument("--sf_vec_size", type=int, default=16) parser.add_argument("--c_dtype", type=cutlass.dtype, default=cutlass.Float16) parser.add_argument("--a_major", choices=["k", "m"], type=str, default="k") @@ -2520,6 +3131,7 @@ if __name__ == "__main__": if len(args.cluster_shape_mn) != 2: parser.error("--cluster_shape_mn must contain exactly 2 values") + # Execute GEMM with appropriate function based on dtype run( args.mnkl, args.ab_dtype, diff --git a/examples/python/CuTeDSL/blackwell/dense_blockscaled_gemm_persistent_amax.py b/examples/python/CuTeDSL/blackwell/dense_blockscaled_gemm_persistent_amax.py new file mode 100644 index 00000000..f423fcce --- /dev/null +++ b/examples/python/CuTeDSL/blackwell/dense_blockscaled_gemm_persistent_amax.py @@ -0,0 +1,2576 @@ +# 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. + +import argparse +from typing import Type, Tuple, Union + +import cuda.bindings.driver as cuda + +import cutlass +import cutlass.cute as cute +from cutlass.cute.nvgpu import cpasync, tcgen05 +from cutlass._mlir.dialects import math +import cutlass.utils as utils +import cutlass.pipeline as pipeline +from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait +import cutlass.utils.blackwell_helpers as sm100_utils +import cutlass.utils.blockscaled_layout as blockscaled_utils +from cutlass.cute.runtime import from_dlpack + +""" +This example provides an experimental implementation of the SM100 batched dense blockscaled GEMM kernel, please note that the APIs and implementation details related to this kernel may change in future releases. + +A high-performance persistent batched dense blockscaled GEMM example for the NVIDIA Blackwell SM100 architecture +using CUTE DSL. +- Matrix A is MxKxL, L is batch dimension, A can be row-major("K") or column-major("M") for MXF8 input type and can only be row-major("K") for MXF4/NVF4 input type +- Matrix B is NxKxL, L is batch dimension, B can be row-major("N") or column-major("K") for MXF8 input type and can only be row-major("K") for MXF4/NVF4 input type +- Matrix C is MxNxL, L is batch dimension, C can be row-major("N") or column-major("M") +- Matrix SFA layout is filled internally according to A shape and BlockScaledBasicChunk, which has M×ceil_div(K, sf_vec_size)×L elements respectively +- Matrix SFB layout is filled internally according to B shape and BlockScaledBasicChunk, which has N×ceil_div(K, sf_vec_size)×L elements respectively + +This GEMM kernel supports the following features: + - Utilizes Tensor Memory Access (TMA) for efficient memory operations + - Utilizes Blackwell's tcgen05.mma for matrix multiply-accumulate (MMA) operations (including 2cta mma instructions) + - Implements TMA multicast with cluster to reduce L2 memory traffic + - Support persistent tile scheduling to better overlap memory load/store with mma between tiles + - Support warp specialization to avoid explicit pipelining between mainloop load and mma + +This GEMM works as follows: +1. DMA warp: Load A and B matrices from global memory (GMEM) to shared memory (SMEM) using TMA operations. +2. MMA warp: + - Load scale factor A/B from shared memory (SMEM) to tensor memory (TMEM) using tcgen05.cp instruction. + - Perform matrix multiply-accumulate (MMA) operations using tcgen05.mma instruction. +3. EPILOGUE warp: + - Load completed accumulator from tensor memory (TMEM) to registers (RMEM) using tcgen05.ld. + - Type convert C matrix to output type. + - Optionally store C matrix from registers (RMEM) to shared memory (SMEM) to global memory (GMEM) with TMA operations, + or directly store C matrix from registers (RMEM) to global memory (GMEM) without TMA operations. + - Optionally accept an elementwise lambda function epilogue_op to apply to the output tensor: + e.g., relu can set epilogue_op = lambda x: cute.where(x > 0, x, cute.full_like(x, 0)) + +SM100 tcgen05.mma.kind.block_scale instructions operate as follows: +- Read matrix A from SMEM +- Read matrix B from SMEM +- Read scalefactor A from TMEM +- Read scalefactor B from TMEM +- Write accumulator to TMEM +The accumulator in TMEM must then be loaded to registers before writing back to GMEM. + +Input arguments to this example is shown below: + +.. code-block:: bash + + python examples/blackwell/dense_blockscaled_gemm_persistent.py \ + --ab_dtype Float4E2M1FN --sf_dtype Float8E8M0FNU --sf_vec_size 16 \ + --c_dtype Float16 \ + --mma_tiler_mn 256,128 --cluster_shape_mn 2,1 \ + --mnkl 8192,8192,1024,1 + +To collect performance with NCU profiler: + +.. code-block:: bash + + ncu python examples/blackwell/dense_blockscaled_gemm_persistent.py \ + --ab_dtype Float4E2M1FN --sf_dtype Float8E8M0FNU --sf_vec_size 16 \ + --c_dtype Float16 \ + --mma_tiler_mn 256,128 --cluster_shape_mn 2,1 \ + --mnkl 8192,8192,1024,1 \ + --warmup_iterations 1 --iterations 10 --skip_ref_check + + +Constraints: +* Supported input data types: mxf8, mxf4, nvf4 + see detailed valid dtype combinations in below Sm100BlockScaledPersistentDenseGemmKernel class documentation +* A/B tensor must have the same data type, mixed data type is not supported (e.g., mxf8 x mxf4) +* Mma tiler M must be 128 or 256(use_2cta_instrs) +* Mma tiler N must be 128 or 256 +* Cluster shape M/N must be positive and power of 2, total cluster size <= 16 +* Cluster shape M must be multiple of 2 if Mma tiler M is 256(use_2cta_instrs) +* The contiguous dimension of A/B/C tensors must be at least 16 bytes aligned, + i.e, number of elements is a multiple of 16 and 32 for Float8 and Float4, respectively. +""" + + +class Sm100BlockScaledPersistentDenseGemmKernel: + """This class implements batched matrix multiplication (C = A x SFA x B x SFB) with support for various data types + and architectural features specific to Blackwell GPUs with persistent tile scheduling and warp specialization. + + :param sf_vec_size: Scalefactor vector size. + :type sf_vec_size: int + :param mma_tiler_mn: Shape of the Matrix Multiply-Accumulate (MMA) tile (M,N) + :type mma_tiler_mn: Tuple[int, int] + :param cluster_shape_mn: Cluster dimensions (M,N) for parallel processing + :type cluster_shape_mn: Tuple[int, int] + + :note: In current version, A and B tensor must have the same data type + - i.e., Float8E4M3FN for A and Float8E5M2 for B is not supported + + :note: Supported combinations of A/B data types, SF data typs and SF vector size: + - MXF8: A/B: Float8E5M2/Float8E4M3FN + SF: Float8E8M0FNU + sf_vec_size: 32 + - MXF4: A/B: Float4E2M1FN + SF: Float8E8M0FNU + sf_vec_size: 32 + - NVF4: A/B: Float4E2M1FN + SF: Float8E8M0FNU/Float8E4M3FN + sf_vec_size: 16 + + :note: Supported accumulator data types: + - Float32 + + :note: Supported C data types: + - Float32 + - Float16/BFloat16 + - Float8E4M3FN/Float8E5M2 + :note: Constraints: + - MMA tiler M must be 128 or 256 (use_2cta_instrs) + - MMA tiler N must be 128/256 + - Cluster shape M must be multiple of 2 if Mma tiler M is 256 + - Cluster shape M/N must be positive and power of 2, total cluster size <= 16 + - Also, Cluster shape M/N must be <= 4 for scale factor multicasts due to limited size of scale factors + + Example: + >>> gemm = Sm100BlockScaledPersistentDenseGemmKernel( + ... sf_vec_size=16, + ... mma_tiler_mn=(256, 128), + ... cluster_shape_mn=(2, 1) + ... ) + >>> gemm(a_tensor, b_tensor, sfa_tensor, sfb_tensor, c_tensor, amax_tensor, max_active_clusters, stream) + """ + + def __init__( + self, + sf_vec_size: int, + mma_tiler_mn: Tuple[int, int], + cluster_shape_mn: Tuple[int, int], + ): + """Initializes the configuration for a Blackwell dense GEMM kernel. + + This configuration includes several key aspects: + + 1. MMA Instruction Settings (tcgen05): + - acc_dtype: Data types for MMA accumulator, always set to Float32 + - sf_vec_size: Scalefactor A/B vector size. + - mma_tiler_mn: The (M, N) shape of the MMA instruction tiler. + + 2. Cluster Shape: + - cluster_shape_mn: The (ClusterM, ClusterN) shape of the CTA cluster. + + :param sf_vec_size: Scalefactor vector size. + :type sf_vec_size: int + :param mma_tiler_mn: Tuple (M, N) shape of the MMA instruction. + :type mma_tiler_mn: Tuple[int, int] + :param cluster_shape_mn: Tuple (ClusterM, ClusterN) shape of the cluster. + :type cluster_shape_mn: Tuple[int, int] + """ + + self.acc_dtype = cutlass.Float32 + self.sf_vec_size = sf_vec_size + self.use_2cta_instrs = mma_tiler_mn[0] == 256 + self.cluster_shape_mn = cluster_shape_mn + # K dimension is deferred in _setup_attributes + self.mma_tiler = (*mma_tiler_mn, 1) + + self.cta_group = ( + tcgen05.CtaGroup.TWO if self.use_2cta_instrs else tcgen05.CtaGroup.ONE + ) + + self.occupancy = 1 + # Set specialized warp ids + self.epilog_warp_id = ( + 0, + 1, + 2, + 3, + ) + self.mma_warp_id = 4 + self.tma_warp_id = 5 + self.threads_per_cta = 32 * len( + (self.mma_warp_id, self.tma_warp_id, *self.epilog_warp_id) + ) + # Set barrier id for cta sync, epilogue sync and tmem ptr sync + self.cta_sync_barrier = pipeline.NamedBarrier( + barrier_id=1, + num_threads=self.threads_per_cta, + ) + self.epilog_sync_barrier = pipeline.NamedBarrier( + barrier_id=2, + num_threads=32 * len(self.epilog_warp_id), + ) + self.tmem_alloc_barrier = pipeline.NamedBarrier( + barrier_id=3, + num_threads=32 * len((self.mma_warp_id, *self.epilog_warp_id)), + ) + + # Amax reduction configuration + self.num_epilog_warps = len(self.epilog_warp_id) + + self.smem_capacity = utils.get_smem_capacity_in_bytes("sm_100") + self.num_tmem_alloc_cols = cute.arch.get_max_tmem_alloc_cols("sm_100") + + def _setup_attributes(self): + """Set up configurations that are dependent on GEMM inputs + + This method configures various attributes based on the input tensor properties + (data types, leading dimensions) and kernel settings: + - Configuring tiled MMA + - Computing MMA/cluster/tile shapes + - Computing cluster layout + - Computing multicast CTAs for A/B/SFA/SFB + - Computing epilogue subtile + - Setting up A/B/SFA/SFB/C stage counts in shared memory + - Computing A/B/SFA/SFB/C shared memory layout + """ + # Compute mma instruction shapes + # (MMA_Tile_Shape_M, MMA_Tile_Shape_N, MMA_Inst_Shape_K) + self.mma_inst_shape_mn = ( + self.mma_tiler[0], + self.mma_tiler[1], + ) + # (CTA_Tile_Shape_M, Round_Up(MMA_Tile_Shape_N, 128), MMA_Inst_Shape_K) + self.mma_inst_shape_mn_sfb = ( + self.mma_inst_shape_mn[0] // (2 if self.use_2cta_instrs else 1), + cute.round_up(self.mma_inst_shape_mn[1], 128), + ) + + tiled_mma = sm100_utils.make_blockscaled_trivial_tiled_mma( + self.a_dtype, + self.a_major_mode, + self.b_major_mode, + self.sf_dtype, + self.sf_vec_size, + self.cta_group, + self.mma_inst_shape_mn, + ) + + tiled_mma_sfb = sm100_utils.make_blockscaled_trivial_tiled_mma( + self.a_dtype, + self.a_major_mode, + self.b_major_mode, + self.sf_dtype, + self.sf_vec_size, + cute.nvgpu.tcgen05.CtaGroup.ONE, + self.mma_inst_shape_mn_sfb, + ) + + # Compute mma/cluster/tile shapes + mma_inst_shape_k = cute.size(tiled_mma.shape_mnk, mode=[2]) + mma_inst_tile_k = 4 + self.mma_tiler = ( + self.mma_inst_shape_mn[0], + self.mma_inst_shape_mn[1], + mma_inst_shape_k * mma_inst_tile_k, + ) + self.mma_tiler_sfb = ( + self.mma_inst_shape_mn_sfb[0], + self.mma_inst_shape_mn_sfb[1], + mma_inst_shape_k * mma_inst_tile_k, + ) + self.cta_tile_shape_mnk = ( + self.mma_tiler[0] // cute.size(tiled_mma.thr_id.shape), + self.mma_tiler[1], + self.mma_tiler[2], + ) + + # Compute cluster layout + self.cluster_layout_vmnk = cute.tiled_divide( + cute.make_layout((*self.cluster_shape_mn, 1)), + (tiled_mma.thr_id.shape,), + ) + self.cluster_layout_sfb_vmnk = cute.tiled_divide( + cute.make_layout((*self.cluster_shape_mn, 1)), + (tiled_mma_sfb.thr_id.shape,), + ) + + # Compute number of multicast CTAs for A/B + self.num_mcast_ctas_a = cute.size(self.cluster_layout_vmnk.shape[2]) + self.num_mcast_ctas_b = cute.size(self.cluster_layout_vmnk.shape[1]) + self.num_mcast_ctas_sfb = cute.size(self.cluster_layout_sfb_vmnk.shape[1]) + self.is_a_mcast = self.num_mcast_ctas_a > 1 + self.is_b_mcast = self.num_mcast_ctas_b > 1 + self.is_sfb_mcast = self.num_mcast_ctas_sfb > 1 + + # Compute epilogue subtile + self.epi_tile = sm100_utils.compute_epilogue_tile_shape( + self.cta_tile_shape_mnk, + self.use_2cta_instrs, + self.c_layout, + self.c_dtype, + ) + + # Setup A/B/C stage count in shared memory and ACC stage count in tensor memory + self.num_acc_stage, self.num_ab_stage, self.num_c_stage = self._compute_stages( + tiled_mma, + self.mma_tiler, + self.a_dtype, + self.b_dtype, + self.epi_tile, + self.c_dtype, + self.c_layout, + self.sf_dtype, + self.sf_vec_size, + self.smem_capacity, + self.occupancy, + ) + + # Compute A/B/SFA/SFB/C shared memory layout + self.a_smem_layout_staged = sm100_utils.make_smem_layout_a( + tiled_mma, + self.mma_tiler, + self.a_dtype, + self.num_ab_stage, + ) + self.b_smem_layout_staged = sm100_utils.make_smem_layout_b( + tiled_mma, + self.mma_tiler, + self.b_dtype, + self.num_ab_stage, + ) + self.sfa_smem_layout_staged = blockscaled_utils.make_smem_layout_sfa( + tiled_mma, + self.mma_tiler, + self.sf_vec_size, + self.num_ab_stage, + ) + self.sfb_smem_layout_staged = blockscaled_utils.make_smem_layout_sfb( + tiled_mma, + self.mma_tiler, + self.sf_vec_size, + self.num_ab_stage, + ) + self.c_smem_layout_staged = sm100_utils.make_smem_layout_epi( + self.c_dtype, + self.c_layout, + self.epi_tile, + self.num_c_stage, + ) + + @cute.jit + def __call__( + self, + a_tensor: cute.Tensor, + b_tensor: cute.Tensor, + sfa_tensor: cute.Tensor, + sfb_tensor: cute.Tensor, + c_tensor: cute.Tensor, + amax_tensor: cute.Tensor, + max_active_clusters: cutlass.Constexpr, + stream: cuda.CUstream, + epilogue_op: cutlass.Constexpr = lambda x: x, + ): + """Execute the GEMM operation in steps: + - Setup static attributes before smem/grid/tma computation + - Setup TMA load/store atoms and tensors + - Compute grid size with regard to hardware constraints + - Define shared storage for kernel + - Launch the kernel synchronously + + :param a_tensor: Input tensor A + :type a_tensor: cute.Tensor + :param b_tensor: Input tensor B + :type b_tensor: cute.Tensor + :param sfa_tensor: Scale factor tensor A + :type sfa_tensor: cute.Tensor + :param sfb_tensor: Scale factor tensor B + :type sfb_tensor: cute.Tensor + :param c_tensor: Output tensor C + :type c_tensor: cute.Tensor + :param amax_tensor: Output tensor for absolute maximum value + :type amax_tensor: cute.Tensor + :param max_active_clusters: Maximum number of active clusters + :type max_active_clusters: cutlass.Constexpr + :param stream: CUDA stream for asynchronous execution + :type stream: cuda.CUstream + :param epilogue_op: Optional elementwise lambda function to apply to the output tensor + :type epilogue_op: cutlass.Constexpr + :raises TypeError: If input data types are incompatible with the MMA instruction. + """ + # Setup static attributes before smem/grid/tma computation + self.a_dtype: Type[cutlass.Numeric] = a_tensor.element_type + self.b_dtype: Type[cutlass.Numeric] = b_tensor.element_type + self.sf_dtype: Type[cutlass.Numeric] = sfa_tensor.element_type + self.c_dtype: Type[cutlass.Numeric] = c_tensor.element_type + self.a_major_mode = utils.LayoutEnum.from_tensor(a_tensor).mma_major_mode() + self.b_major_mode = utils.LayoutEnum.from_tensor(b_tensor).mma_major_mode() + self.c_layout = utils.LayoutEnum.from_tensor(c_tensor) + + # Check if input data types are compatible with MMA instruction + if cutlass.const_expr(self.a_dtype != self.b_dtype): + raise TypeError(f"Type must match: {self.a_dtype} != {self.b_dtype}") + + # Setup attributes that dependent on gemm inputs + self._setup_attributes() + + # Setup sfa/sfb tensor by filling A/B tensor to scale factor atom layout + # ((Atom_M, Rest_M),(Atom_K, Rest_K),RestL) + sfa_layout = blockscaled_utils.tile_atom_to_shape_SF( + a_tensor.shape, self.sf_vec_size + ) + sfa_tensor = cute.make_tensor(sfa_tensor.iterator, sfa_layout) + + # ((Atom_N, Rest_N),(Atom_K, Rest_K),RestL) + sfb_layout = blockscaled_utils.tile_atom_to_shape_SF( + b_tensor.shape, self.sf_vec_size + ) + sfb_tensor = cute.make_tensor(sfb_tensor.iterator, sfb_layout) + + tiled_mma = sm100_utils.make_blockscaled_trivial_tiled_mma( + self.a_dtype, + self.a_major_mode, + self.b_major_mode, + self.sf_dtype, + self.sf_vec_size, + self.cta_group, + self.mma_inst_shape_mn, + ) + + tiled_mma_sfb = sm100_utils.make_blockscaled_trivial_tiled_mma( + self.a_dtype, + self.a_major_mode, + self.b_major_mode, + self.sf_dtype, + self.sf_vec_size, + cute.nvgpu.tcgen05.CtaGroup.ONE, + self.mma_inst_shape_mn_sfb, + ) + atom_thr_size = cute.size(tiled_mma.thr_id.shape) + + # Setup TMA load for A + a_op = sm100_utils.cluster_shape_to_tma_atom_A( + self.cluster_shape_mn, tiled_mma.thr_id + ) + a_smem_layout = cute.slice_(self.a_smem_layout_staged, (None, None, None, 0)) + tma_atom_a, tma_tensor_a = cute.nvgpu.make_tiled_tma_atom_A( + a_op, + a_tensor, + a_smem_layout, + self.mma_tiler, + tiled_mma, + self.cluster_layout_vmnk.shape, + ) + + # Setup TMA load for B + b_op = sm100_utils.cluster_shape_to_tma_atom_B( + self.cluster_shape_mn, tiled_mma.thr_id + ) + b_smem_layout = cute.slice_(self.b_smem_layout_staged, (None, None, None, 0)) + tma_atom_b, tma_tensor_b = cute.nvgpu.make_tiled_tma_atom_B( + b_op, + b_tensor, + b_smem_layout, + self.mma_tiler, + tiled_mma, + self.cluster_layout_vmnk.shape, + ) + + # Setup TMA load for SFA + sfa_op = sm100_utils.cluster_shape_to_tma_atom_A( + self.cluster_shape_mn, tiled_mma.thr_id + ) + sfa_smem_layout = cute.slice_( + self.sfa_smem_layout_staged, (None, None, None, 0) + ) + tma_atom_sfa, tma_tensor_sfa = cute.nvgpu.make_tiled_tma_atom_A( + sfa_op, + sfa_tensor, + sfa_smem_layout, + self.mma_tiler, + tiled_mma, + self.cluster_layout_vmnk.shape, + internal_type=cutlass.Int16, + ) + + # Setup TMA load for SFB + sfb_op = sm100_utils.cluster_shape_to_tma_atom_SFB( + self.cluster_shape_mn, tiled_mma.thr_id + ) + sfb_smem_layout = cute.slice_( + self.sfb_smem_layout_staged, (None, None, None, 0) + ) + tma_atom_sfb, tma_tensor_sfb = cute.nvgpu.make_tiled_tma_atom_B( + sfb_op, + sfb_tensor, + sfb_smem_layout, + self.mma_tiler_sfb, + tiled_mma_sfb, + self.cluster_layout_sfb_vmnk.shape, + internal_type=cutlass.Int16, + ) + + a_copy_size = cute.size_in_bytes(self.a_dtype, a_smem_layout) + b_copy_size = cute.size_in_bytes(self.b_dtype, b_smem_layout) + sfa_copy_size = cute.size_in_bytes(self.sf_dtype, sfa_smem_layout) + sfb_copy_size = cute.size_in_bytes(self.sf_dtype, sfb_smem_layout) + self.num_tma_load_bytes = ( + a_copy_size + b_copy_size + sfa_copy_size + sfb_copy_size + ) * atom_thr_size + + # Setup TMA store for C + epi_smem_layout = cute.slice_(self.c_smem_layout_staged, (None, None, 0)) + tma_atom_c, tma_tensor_c = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileS2GOp(), + c_tensor, + epi_smem_layout, + self.epi_tile, + ) + + # Compute grid size + self.tile_sched_params, grid = self._compute_grid( + c_tensor, + self.cta_tile_shape_mnk, + self.cluster_shape_mn, + max_active_clusters, + ) + + self.buffer_align_bytes = 1024 + + # Define shared storage for kernel + @cute.struct + class SharedStorage: + ab_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage] + ab_empty_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage] + acc_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage] + acc_empty_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage] + tmem_dealloc_mbar_ptr: cutlass.Int64 + tmem_holding_buf: cutlass.Int32 + # (EPI_TILE_M, EPI_TILE_N, STAGE) + sC: cute.struct.Align[ + cute.struct.MemRange[ + self.c_dtype, + cute.cosize(self.c_smem_layout_staged.outer), + ], + self.buffer_align_bytes, + ] + # (MMA, MMA_M, MMA_K, STAGE) + sA: cute.struct.Align[ + cute.struct.MemRange[ + self.a_dtype, cute.cosize(self.a_smem_layout_staged.outer) + ], + self.buffer_align_bytes, + ] + # (MMA, MMA_N, MMA_K, STAGE) + sB: cute.struct.Align[ + cute.struct.MemRange[ + self.b_dtype, cute.cosize(self.b_smem_layout_staged.outer) + ], + self.buffer_align_bytes, + ] + # (MMA, MMA_M, MMA_K, STAGE) + sSFA: cute.struct.Align[ + cute.struct.MemRange[ + self.sf_dtype, cute.cosize(self.sfa_smem_layout_staged) + ], + self.buffer_align_bytes, + ] + # (MMA, MMA_N, MMA_K, STAGE) + sSFB: cute.struct.Align[ + cute.struct.MemRange[ + self.sf_dtype, cute.cosize(self.sfb_smem_layout_staged) + ], + self.buffer_align_bytes, + ] + # Amax reduction shared memory (one FP32 per epilogue warp) + # Use smaller alignment for amax since it's only 16 bytes + sAmax: cute.struct.Align[ + cute.struct.MemRange[cutlass.Float32, self.num_epilog_warps], + 16, + ] + + self.shared_storage = SharedStorage + # Launch the kernel synchronously + self.kernel( + tiled_mma, + tiled_mma_sfb, + tma_atom_a, + tma_tensor_a, + tma_atom_b, + tma_tensor_b, + tma_atom_sfa, + tma_tensor_sfa, + tma_atom_sfb, + tma_tensor_sfb, + tma_atom_c, + tma_tensor_c, + amax_tensor, + self.cluster_layout_vmnk, + self.cluster_layout_sfb_vmnk, + self.a_smem_layout_staged, + self.b_smem_layout_staged, + self.sfa_smem_layout_staged, + self.sfb_smem_layout_staged, + self.c_smem_layout_staged, + self.epi_tile, + self.tile_sched_params, + epilogue_op, + ).launch( + grid=grid, + block=[self.threads_per_cta, 1, 1], + cluster=(*self.cluster_shape_mn, 1), + stream=stream, + min_blocks_per_mp=1, + ) + return + + # GPU device kernel + @cute.kernel + def kernel( + self, + tiled_mma: cute.TiledMma, + tiled_mma_sfb: cute.TiledMma, + tma_atom_a: cute.CopyAtom, + mA_mkl: cute.Tensor, + tma_atom_b: cute.CopyAtom, + mB_nkl: cute.Tensor, + tma_atom_sfa: cute.CopyAtom, + mSFA_mkl: cute.Tensor, + tma_atom_sfb: cute.CopyAtom, + mSFB_nkl: cute.Tensor, + tma_atom_c: cute.CopyAtom, + mC_mnl: cute.Tensor, + mAmax: cute.Tensor, + cluster_layout_vmnk: cute.Layout, + cluster_layout_sfb_vmnk: cute.Layout, + a_smem_layout_staged: cute.ComposedLayout, + b_smem_layout_staged: cute.ComposedLayout, + sfa_smem_layout_staged: cute.Layout, + sfb_smem_layout_staged: cute.Layout, + c_smem_layout_staged: Union[cute.Layout, cute.ComposedLayout], + epi_tile: cute.Tile, + tile_sched_params: utils.PersistentTileSchedulerParams, + epilogue_op: cutlass.Constexpr, + ): + """ + GPU device kernel performing the Persistent batched GEMM computation. + """ + warp_idx = cute.arch.warp_idx() + warp_idx = cute.arch.make_warp_uniform(warp_idx) + + # + # Prefetch tma desc + # + if warp_idx == self.tma_warp_id: + cpasync.prefetch_descriptor(tma_atom_a) + cpasync.prefetch_descriptor(tma_atom_b) + cpasync.prefetch_descriptor(tma_atom_sfa) + cpasync.prefetch_descriptor(tma_atom_sfb) + cpasync.prefetch_descriptor(tma_atom_c) + + use_2cta_instrs = cute.size(tiled_mma.thr_id.shape) == 2 + + # + # Setup cta/thread coordinates + # + # Coords inside cluster + bidx, bidy, bidz = cute.arch.block_idx() + mma_tile_coord_v = bidx % cute.size(tiled_mma.thr_id.shape) + is_leader_cta = mma_tile_coord_v == 0 + cta_rank_in_cluster = cute.arch.make_warp_uniform( + cute.arch.block_idx_in_cluster() + ) + block_in_cluster_coord_vmnk = cluster_layout_vmnk.get_flat_coord( + cta_rank_in_cluster + ) + block_in_cluster_coord_sfb_vmnk = cluster_layout_sfb_vmnk.get_flat_coord( + cta_rank_in_cluster + ) + # Coord inside cta + tidx, _, _ = cute.arch.thread_idx() + + # + # Alloc and init: a+b full/empty, accumulator full/empty, tensor memory dealloc barrier + # + smem = utils.SmemAllocator() + storage = smem.allocate(self.shared_storage) + + # Initialize mainloop ab_pipeline (barrier) and states + ab_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) + num_tma_producer = self.num_mcast_ctas_a + self.num_mcast_ctas_b - 1 + ab_pipeline_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, num_tma_producer + ) + ab_pipeline = pipeline.PipelineTmaUmma.create( + barrier_storage=storage.ab_full_mbar_ptr.data_ptr(), + num_stages=self.num_ab_stage, + producer_group=ab_pipeline_producer_group, + consumer_group=ab_pipeline_consumer_group, + tx_count=self.num_tma_load_bytes, + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ) + + # Initialize acc_pipeline (barrier) and states + acc_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) + num_acc_consumer_threads = len(self.epilog_warp_id) * ( + 2 if use_2cta_instrs else 1 + ) + acc_pipeline_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, num_acc_consumer_threads + ) + acc_pipeline = pipeline.PipelineUmmaAsync.create( + barrier_storage=storage.acc_full_mbar_ptr.data_ptr(), + num_stages=self.num_acc_stage, + producer_group=acc_pipeline_producer_group, + consumer_group=acc_pipeline_consumer_group, + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ) + + # Tensor memory dealloc barrier init + tmem = utils.TmemAllocator( + storage.tmem_holding_buf, + barrier_for_retrieve=self.tmem_alloc_barrier, + allocator_warp_id=self.epilog_warp_id[0], + is_two_cta=use_2cta_instrs, + two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar_ptr, + ) + + # Cluster arrive after barrier init + pipeline_init_arrive(cluster_shape_mn=self.cluster_shape_mn, is_relaxed=True) + + # + # Setup smem tensor A/B/SFA/SFB/C + # + # (EPI_TILE_M, EPI_TILE_N, STAGE) + sC = storage.sC.get_tensor( + c_smem_layout_staged.outer, swizzle=c_smem_layout_staged.inner + ) + # (MMA, MMA_M, MMA_K, STAGE) + sA = storage.sA.get_tensor( + a_smem_layout_staged.outer, swizzle=a_smem_layout_staged.inner + ) + # (MMA, MMA_N, MMA_K, STAGE) + sB = storage.sB.get_tensor( + b_smem_layout_staged.outer, swizzle=b_smem_layout_staged.inner + ) + # (MMA, MMA_M, MMA_K, STAGE) + sSFA = storage.sSFA.get_tensor(sfa_smem_layout_staged) + # (MMA, MMA_N, MMA_K, STAGE) + sSFB = storage.sSFB.get_tensor(sfb_smem_layout_staged) + + # Shared memory for amax reduction (one FP32 per epilogue warp) + # Simple 1D layout + amax_layout = cute.make_layout((self.num_epilog_warps,)) + sAmax = storage.sAmax.get_tensor(amax_layout) + + # + # Compute multicast mask for A/B/SFA/SFB buffer full + # + a_full_mcast_mask = None + b_full_mcast_mask = None + sfa_full_mcast_mask = None + sfb_full_mcast_mask = None + if cutlass.const_expr(self.is_a_mcast or self.is_b_mcast or use_2cta_instrs): + a_full_mcast_mask = cpasync.create_tma_multicast_mask( + cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=2 + ) + b_full_mcast_mask = cpasync.create_tma_multicast_mask( + cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=1 + ) + sfa_full_mcast_mask = cpasync.create_tma_multicast_mask( + cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=2 + ) + sfb_full_mcast_mask = cpasync.create_tma_multicast_mask( + cluster_layout_sfb_vmnk, block_in_cluster_coord_sfb_vmnk, mcast_mode=1 + ) + + # + # Local_tile partition global tensors + # + # (bM, bK, RestM, RestK, RestL) + gA_mkl = cute.local_tile( + mA_mkl, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None) + ) + # (bN, bK, RestN, RestK, RestL) + gB_nkl = cute.local_tile( + mB_nkl, cute.slice_(self.mma_tiler, (0, None, None)), (None, None, None) + ) + # (bM, bK, RestM, RestK, RestL) + gSFA_mkl = cute.local_tile( + mSFA_mkl, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None) + ) + # (bN, bK, RestN, RestK, RestL) + gSFB_nkl = cute.local_tile( + mSFB_nkl, cute.slice_(self.mma_tiler, (0, None, None)), (None, None, None) + ) + # (bM, bN, RestM, RestN, RestL) + gC_mnl = cute.local_tile( + mC_mnl, cute.slice_(self.mma_tiler, (None, None, 0)), (None, None, None) + ) + k_tile_cnt = cute.size(gA_mkl, mode=[3]) + + # + # Partition global tensor for TiledMMA_A/B/C + # + thr_mma = tiled_mma.get_slice(mma_tile_coord_v) + thr_mma_sfb = tiled_mma_sfb.get_slice(mma_tile_coord_v) + # (MMA, MMA_M, MMA_K, RestM, RestK, RestL) + tCgA = thr_mma.partition_A(gA_mkl) + # (MMA, MMA_N, MMA_K, RestN, RestK, RestL) + tCgB = thr_mma.partition_B(gB_nkl) + # (MMA, MMA_M, MMA_K, RestM, RestK, RestL) + tCgSFA = thr_mma.partition_A(gSFA_mkl) + # (MMA, MMA_N, MMA_K, RestN, RestK, RestL) + tCgSFB = thr_mma_sfb.partition_B(gSFB_nkl) + # (MMA, MMA_M, MMA_N, RestM, RestN, RestL) + tCgC = thr_mma.partition_C(gC_mnl) + + # + # Partition global/shared tensor for TMA load A/B + # + # TMA load A partition_S/D + a_cta_layout = cute.make_layout( + cute.slice_(cluster_layout_vmnk, (0, 0, None, 0)).shape + ) + # ((atom_v, rest_v), STAGE) + # ((atom_v, rest_v), RestM, RestK, RestL) + tAsA, tAgA = cpasync.tma_partition( + tma_atom_a, + block_in_cluster_coord_vmnk[2], + a_cta_layout, + cute.group_modes(sA, 0, 3), + cute.group_modes(tCgA, 0, 3), + ) + # TMA load B partition_S/D + b_cta_layout = cute.make_layout( + cute.slice_(cluster_layout_vmnk, (0, None, 0, 0)).shape + ) + # ((atom_v, rest_v), STAGE) + # ((atom_v, rest_v), RestN, RestK, RestL) + tBsB, tBgB = cpasync.tma_partition( + tma_atom_b, + block_in_cluster_coord_vmnk[1], + b_cta_layout, + cute.group_modes(sB, 0, 3), + cute.group_modes(tCgB, 0, 3), + ) + + # TMA load scaled factor A partition_S/D + sfa_cta_layout = a_cta_layout + # ((atom_v, rest_v), STAGE) + # ((atom_v, rest_v), RestM, RestK, RestL) + tAsSFA, tAgSFA = cute.nvgpu.cpasync.tma_partition( + tma_atom_sfa, + block_in_cluster_coord_vmnk[2], + sfa_cta_layout, + cute.group_modes(sSFA, 0, 3), + cute.group_modes(tCgSFA, 0, 3), + ) + tAsSFA = cute.filter_zeros(tAsSFA) + tAgSFA = cute.filter_zeros(tAgSFA) + + # TMA load scaled factor B partition_S/D + sfb_cta_layout = cute.make_layout( + cute.slice_(cluster_layout_sfb_vmnk, (0, None, 0, 0)).shape + ) + # ((atom_v, rest_v), STAGE) + # ((atom_v, rest_v), RestN, RestK, RestL) + tBsSFB, tBgSFB = cute.nvgpu.cpasync.tma_partition( + tma_atom_sfb, + block_in_cluster_coord_sfb_vmnk[1], + sfb_cta_layout, + cute.group_modes(sSFB, 0, 3), + cute.group_modes(tCgSFB, 0, 3), + ) + tBsSFB = cute.filter_zeros(tBsSFB) + tBgSFB = cute.filter_zeros(tBgSFB) + + # + # Partition shared/tensor memory tensor for TiledMMA_A/B/C + # + # (MMA, MMA_M, MMA_K, STAGE) + tCrA = tiled_mma.make_fragment_A(sA) + # (MMA, MMA_N, MMA_K, STAGE) + tCrB = tiled_mma.make_fragment_B(sB) + # (MMA, MMA_M, MMA_N) + acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2]) + # (MMA, MMA_M, MMA_N, STAGE) + tCtAcc_fake = tiled_mma.make_fragment_C( + cute.append(acc_shape, self.num_acc_stage) + ) + + # + # Cluster wait before tensor memory alloc + # + pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mn) + + # + # Specialized TMA load warp + # + if warp_idx == self.tma_warp_id: + # + # Persistent tile scheduling loop + # + tile_sched = utils.StaticPersistentTileScheduler.create( + tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim() + ) + work_tile = tile_sched.initial_work_tile_info() + + ab_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.num_ab_stage + ) + + while work_tile.is_valid_tile: + # Get tile coord from tile scheduler + cur_tile_coord = work_tile.tile_idx + mma_tile_coord_mnl = ( + cur_tile_coord[0] // cute.size(tiled_mma.thr_id.shape), + cur_tile_coord[1], + cur_tile_coord[2], + ) + + # + # Slice to per mma tile index + # + # ((atom_v, rest_v), RestK) + tAgA_slice = tAgA[ + (None, mma_tile_coord_mnl[0], None, mma_tile_coord_mnl[2]) + ] + # ((atom_v, rest_v), RestK) + tBgB_slice = tBgB[ + (None, mma_tile_coord_mnl[1], None, mma_tile_coord_mnl[2]) + ] + + # ((atom_v, rest_v), RestK) + tAgSFA_slice = tAgSFA[ + (None, mma_tile_coord_mnl[0], None, mma_tile_coord_mnl[2]) + ] + # ((atom_v, rest_v), RestK) + tBgSFB_slice = tBgSFB[ + (None, mma_tile_coord_mnl[1], None, mma_tile_coord_mnl[2]) + ] + + # Peek (try_wait) AB buffer empty for k_tile = prefetch_k_tile_cnt + ab_producer_state.reset_count() + peek_ab_empty_status = cutlass.Boolean(1) + if ab_producer_state.count < k_tile_cnt: + peek_ab_empty_status = ab_pipeline.producer_try_acquire( + ab_producer_state + ) + # + # Tma load loop + # + for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1): + # Conditionally wait for AB buffer empty + ab_pipeline.producer_acquire( + ab_producer_state, peek_ab_empty_status + ) + + # TMA load A/B/SFA/SFB + cute.copy( + tma_atom_a, + tAgA_slice[(None, ab_producer_state.count)], + tAsA[(None, ab_producer_state.index)], + tma_bar_ptr=ab_pipeline.producer_get_barrier(ab_producer_state), + mcast_mask=a_full_mcast_mask, + ) + cute.copy( + tma_atom_b, + tBgB_slice[(None, ab_producer_state.count)], + tBsB[(None, ab_producer_state.index)], + tma_bar_ptr=ab_pipeline.producer_get_barrier(ab_producer_state), + mcast_mask=b_full_mcast_mask, + ) + cute.copy( + tma_atom_sfa, + tAgSFA_slice[(None, ab_producer_state.count)], + tAsSFA[(None, ab_producer_state.index)], + tma_bar_ptr=ab_pipeline.producer_get_barrier(ab_producer_state), + mcast_mask=sfa_full_mcast_mask, + ) + cute.copy( + tma_atom_sfb, + tBgSFB_slice[(None, ab_producer_state.count)], + tBsSFB[(None, ab_producer_state.index)], + tma_bar_ptr=ab_pipeline.producer_get_barrier(ab_producer_state), + mcast_mask=sfb_full_mcast_mask, + ) + + # Peek (try_wait) AB buffer empty for k_tile = prefetch_k_tile_cnt + k_tile + 1 + ab_producer_state.advance() + peek_ab_empty_status = cutlass.Boolean(1) + if ab_producer_state.count < k_tile_cnt: + peek_ab_empty_status = ab_pipeline.producer_try_acquire( + ab_producer_state + ) + + # + # Advance to next tile + # + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + + # + # Wait A/B buffer empty + # + ab_pipeline.producer_tail(ab_producer_state) + + # + # Specialized MMA warp + # + if warp_idx == self.mma_warp_id: + # + # Bar sync for retrieve tensor memory ptr from shared mem + # + tmem.wait_for_alloc() + + # + # Retrieving tensor memory ptr and make accumulator/SFA/SFB tensor + # + acc_tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + # Make accumulator tmem tensor + # (MMA, MMA_M, MMA_N, STAGE) + tCtAcc_base = cute.make_tensor(acc_tmem_ptr, tCtAcc_fake.layout) + + # Make SFA tmem tensor + sfa_tmem_ptr = cute.recast_ptr( + acc_tmem_ptr + tcgen05.find_tmem_tensor_col_offset(tCtAcc_base), + dtype=self.sf_dtype, + ) + # (MMA, MMA_M, MMA_K) + tCtSFA_layout = blockscaled_utils.make_tmem_layout_sfa( + tiled_mma, + self.mma_tiler, + self.sf_vec_size, + cute.slice_(sfa_smem_layout_staged, (None, None, None, 0)), + ) + tCtSFA = cute.make_tensor(sfa_tmem_ptr, tCtSFA_layout) + + # Make SFB tmem tensor + sfb_tmem_ptr = cute.recast_ptr( + acc_tmem_ptr + + tcgen05.find_tmem_tensor_col_offset(tCtAcc_base) + + tcgen05.find_tmem_tensor_col_offset(tCtSFA), + dtype=self.sf_dtype, + ) + # (MMA, MMA_N, MMA_K) + tCtSFB_layout = blockscaled_utils.make_tmem_layout_sfb( + tiled_mma, + self.mma_tiler, + self.sf_vec_size, + cute.slice_(sfb_smem_layout_staged, (None, None, None, 0)), + ) + tCtSFB = cute.make_tensor(sfb_tmem_ptr, tCtSFB_layout) + # + # Partition for S2T copy of SFA/SFB + # + ( + tiled_copy_s2t_sfa, + tCsSFA_compact_s2t, + tCtSFA_compact_s2t, + ) = self.mainloop_s2t_copy_and_partition(sSFA, tCtSFA) + ( + tiled_copy_s2t_sfb, + tCsSFB_compact_s2t, + tCtSFB_compact_s2t, + ) = self.mainloop_s2t_copy_and_partition(sSFB, tCtSFB) + + # + # Persistent tile scheduling loop + # + tile_sched = utils.StaticPersistentTileScheduler.create( + tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim() + ) + work_tile = tile_sched.initial_work_tile_info() + + ab_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_ab_stage + ) + acc_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.num_acc_stage + ) + + while work_tile.is_valid_tile: + # Get tile coord from tile scheduler + cur_tile_coord = work_tile.tile_idx + mma_tile_coord_mnl = ( + cur_tile_coord[0] // cute.size(tiled_mma.thr_id.shape), + cur_tile_coord[1], + cur_tile_coord[2], + ) + + # Set tensor memory buffer for current tile + # (MMA, MMA_M, MMA_N) + tCtAcc = tCtAcc_base[(None, None, None, acc_producer_state.index)] + + # Peek (try_wait) AB buffer full for k_tile = 0 + ab_consumer_state.reset_count() + peek_ab_full_status = cutlass.Boolean(1) + if ab_consumer_state.count < k_tile_cnt and is_leader_cta: + peek_ab_full_status = ab_pipeline.consumer_try_wait( + ab_consumer_state + ) + + # + # Wait for accumulator buffer empty + # + if is_leader_cta: + acc_pipeline.producer_acquire(acc_producer_state) + + # + # Reset the ACCUMULATE field for each tile + # + tiled_mma.set(tcgen05.Field.ACCUMULATE, False) + + # + # Mma mainloop + # + for k_tile in range(k_tile_cnt): + if is_leader_cta: + # Conditionally wait for AB buffer full + ab_pipeline.consumer_wait( + ab_consumer_state, peek_ab_full_status + ) + + # Copy SFA/SFB from smem to tmem + s2t_stage_coord = ( + None, + None, + None, + None, + ab_consumer_state.index, + ) + tCsSFA_compact_s2t_staged = tCsSFA_compact_s2t[s2t_stage_coord] + tCsSFB_compact_s2t_staged = tCsSFB_compact_s2t[s2t_stage_coord] + cute.copy( + tiled_copy_s2t_sfa, + tCsSFA_compact_s2t_staged, + tCtSFA_compact_s2t, + ) + cute.copy( + tiled_copy_s2t_sfb, + tCsSFB_compact_s2t_staged, + tCtSFB_compact_s2t, + ) + + # tCtAcc += tCrA * tCrSFA * tCrB * tCrSFB + num_kblocks = cute.size(tCrA, mode=[2]) + for kblock_idx in cutlass.range(num_kblocks, unroll_full=True): + kblock_coord = ( + None, + None, + kblock_idx, + ab_consumer_state.index, + ) + + # Set SFA/SFB tensor to tiled_mma + sf_kblock_coord = (None, None, kblock_idx) + tiled_mma.set( + tcgen05.Field.SFA, + tCtSFA[sf_kblock_coord].iterator, + ) + tiled_mma.set( + tcgen05.Field.SFB, + tCtSFB[sf_kblock_coord].iterator, + ) + + cute.gemm( + tiled_mma, + tCtAcc, + tCrA[kblock_coord], + tCrB[kblock_coord], + tCtAcc, + ) + + # Enable accumulate on tCtAcc after first kblock + tiled_mma.set(tcgen05.Field.ACCUMULATE, True) + + # Async arrive AB buffer empty + ab_pipeline.consumer_release(ab_consumer_state) + + # Peek (try_wait) AB buffer full for k_tile = k_tile + 1 + ab_consumer_state.advance() + peek_ab_full_status = cutlass.Boolean(1) + if ab_consumer_state.count < k_tile_cnt: + if is_leader_cta: + peek_ab_full_status = ab_pipeline.consumer_try_wait( + ab_consumer_state + ) + + # + # Async arrive accumulator buffer full + # + if is_leader_cta: + acc_pipeline.producer_commit(acc_producer_state) + acc_producer_state.advance() + + # + # Advance to next tile + # + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + + # + # Wait for accumulator buffer empty + # + acc_pipeline.producer_tail(acc_producer_state) + # + # Specialized epilogue warps + # + if warp_idx < self.mma_warp_id: + # + # Alloc tensor memory buffer + # + tmem.allocate(self.num_tmem_alloc_cols) + + # + # Bar sync for retrieve tensor memory ptr from shared memory + # + tmem.wait_for_alloc() + + # + # Retrieving tensor memory ptr and make accumulator tensor + # + acc_tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + # (MMA, MMA_M, MMA_N, STAGE) + tCtAcc_base = cute.make_tensor(acc_tmem_ptr, tCtAcc_fake.layout) + + # + # Partition for epilogue + # + epi_tidx = tidx + ( + tiled_copy_t2r, + tTR_tAcc_base, + tTR_rAcc, + ) = self.epilog_tmem_copy_and_partition( + epi_tidx, tCtAcc_base, tCgC, epi_tile, use_2cta_instrs + ) + + tTR_rC = cute.make_rmem_tensor(tTR_rAcc.shape, self.c_dtype) + tiled_copy_r2s, tRS_rC, tRS_sC = self.epilog_smem_copy_and_partition( + tiled_copy_t2r, tTR_rC, epi_tidx, sC + ) + ( + tma_atom_c, + bSG_sC, + bSG_gC_partitioned, + ) = self.epilog_gmem_copy_and_partition( + epi_tidx, tma_atom_c, tCgC, epi_tile, sC + ) + + # + # Persistent tile scheduling loop + # + tile_sched = utils.StaticPersistentTileScheduler.create( + tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim() + ) + work_tile = tile_sched.initial_work_tile_info() + + acc_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_acc_stage + ) + + # Threads/warps participating in tma store pipeline + c_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + 32 * len(self.epilog_warp_id), + ) + c_pipeline = pipeline.PipelineTmaStore.create( + num_stages=self.num_c_stage, + producer_group=c_producer_group, + ) + + while work_tile.is_valid_tile: + # Get tile coord from tile scheduler + cur_tile_coord = work_tile.tile_idx + mma_tile_coord_mnl = ( + cur_tile_coord[0] // cute.size(tiled_mma.thr_id.shape), + cur_tile_coord[1], + cur_tile_coord[2], + ) + + # + # Slice to per mma tile index + # + # ((ATOM_V, REST_V), EPI_M, EPI_N) + bSG_gC = bSG_gC_partitioned[ + ( + None, + None, + None, + *mma_tile_coord_mnl, + ) + ] + + # Set tensor memory buffer for current tile + # (T2R, T2R_M, T2R_N, EPI_M, EPI_M) + tTR_tAcc = tTR_tAcc_base[ + (None, None, None, None, None, acc_consumer_state.index) + ] + + # + # Wait for accumulator buffer full + # + acc_pipeline.consumer_wait(acc_consumer_state) + + tTR_tAcc = cute.group_modes(tTR_tAcc, 3, cute.rank(tTR_tAcc)) + bSG_gC = cute.group_modes(bSG_gC, 1, cute.rank(bSG_gC)) + + # + # Store accumulator to global memory in subtiles + # + subtile_cnt = cute.size(tTR_tAcc.shape, mode=[3]) + num_prev_subtiles = tile_sched.num_tiles_executed * subtile_cnt + + # Initialize thread-local amax accumulator for this tile + # Use 0.0 as initial value since we're computing absolute maximum + thread_tile_amax = cutlass.Float32(0.0) + + for subtile_idx in cutlass.range(subtile_cnt): + # + # Load accumulator from tensor memory buffer to register + # + tTR_tAcc_mn = tTR_tAcc[(None, None, None, subtile_idx)] + cute.copy(tiled_copy_t2r, tTR_tAcc_mn, tTR_rAcc) + + # Accumulate thread-level amax across all subtiles in this tile + # Note: We need absolute value maximum, so take abs first + acc_values = tTR_rAcc.load() + # Apply element-wise absolute value using math.absf (supports vectors) + abs_acc_values_ir = math.absf( + acc_values.ir_value() # operand (positional) + ) + abs_acc_values = type(acc_values)( + abs_acc_values_ir, acc_values.shape, acc_values.dtype + ) + subtile_amax = abs_acc_values.reduce( + cute.ReductionOp.MAX, + cutlass.Float32(0.0), + 0, # Use 0.0 as init for abs values + ) + thread_tile_amax = cute.arch.fmax(thread_tile_amax, subtile_amax) + + # + # Convert to C type + # + acc_vec = tiled_copy_r2s.retile(tTR_rAcc).load() + acc_vec = epilogue_op(acc_vec.to(self.c_dtype)) + tRS_rC.store(acc_vec) + + # + # Store C to shared memory + # + c_buffer = (num_prev_subtiles + subtile_idx) % self.num_c_stage + cute.copy( + tiled_copy_r2s, + tRS_rC, + tRS_sC[(None, None, None, c_buffer)], + ) + # Fence and barrier to make sure shared memory store is visible to TMA store + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) + self.epilog_sync_barrier.arrive_and_wait() + + # + # TMA store C to global memory + # + if warp_idx == self.epilog_warp_id[0]: + cute.copy( + tma_atom_c, + bSG_sC[(None, c_buffer)], + bSG_gC[(None, subtile_idx)], + ) + # Fence and barrier to make sure shared memory store is visible to TMA store + c_pipeline.producer_commit() + c_pipeline.producer_acquire() + self.epilog_sync_barrier.arrive_and_wait() + + # Perform amax reduction after all subtiles are processed + # Warp-level reduction using wrapper function + warp_amax = cute.arch.warp_redux_sync( + value=thread_tile_amax, + kind="fmax", + mask_and_clamp=0xFFFFFFFF, + nan=True, + ) + # Each epilogue warp's lane 0 writes warp amax to shared memory + if cute.arch.lane_idx() == 0: + sAmax[warp_idx] = cutlass.Float32(warp_amax) + + # Ensure all epilogue warps complete their writes before block reduction + self.epilog_sync_barrier.arrive_and_wait() + + # Block-level reduction: only first epilogue warp's lane 0 handles this + if warp_idx == self.epilog_warp_id[0] and cute.arch.lane_idx() == 0: + block_amax = cutlass.Float32( + 0.0 + ) # Initial value for absolute maximum + for i in cutlass.range(self.num_epilog_warps): + warp_amax_val = sAmax[i] + block_amax = cute.arch.fmax(block_amax, warp_amax_val) + + # Global atomic max (accumulates across all tiles for final tensor amax) + # Since we compute absolute values, all values are non-negative + # Use wrapper function for atomic max operation + _ = cute.arch.atomic_max_float32( + ptr=mAmax.iterator.llvm_ptr, value=block_amax + ) + # + # Async arrive accumulator buffer empty + # + with cute.arch.elect_one(): + acc_pipeline.consumer_release(acc_consumer_state) + acc_consumer_state.advance() + + # + # Advance to next tile + # + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + + # + # Dealloc the tensor memory buffer + # + tmem.relinquish_alloc_permit() + self.epilog_sync_barrier.arrive_and_wait() + tmem.free(acc_tmem_ptr) + # + # Wait for C store complete + # + c_pipeline.producer_tail() + + def mainloop_s2t_copy_and_partition( + self, + sSF: cute.Tensor, + tSF: cute.Tensor, + ) -> Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]: + """ + Make tiledCopy for smem to tmem load for scale factor tensor, then use it to partition smem memory (source) and tensor memory (destination). + + :param sSF: The scale factor tensor in smem + :type sSF: cute.Tensor + :param tSF: The scale factor tensor in tmem + :type tSF: cute.Tensor + + :return: A tuple containing (tiled_copy_s2t, tCsSF_compact_s2t, tCtSF_compact_s2t) where: + - tiled_copy_s2t: The tiled copy operation for smem to tmem load for scale factor tensor(s2t) + - tCsSF_compact_s2t: The partitioned scale factor tensor in smem + - tSF_compact_s2t: The partitioned scale factor tensor in tmem + :rtype: Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor] + """ + # (MMA, MMA_MN, MMA_K, STAGE) + tCsSF_compact = cute.filter_zeros(sSF) + # (MMA, MMA_MN, MMA_K) + tCtSF_compact = cute.filter_zeros(tSF) + + # Make S2T CopyAtom and tiledCopy + copy_atom_s2t = cute.make_copy_atom( + tcgen05.Cp4x32x128bOp(self.cta_group), + self.sf_dtype, + ) + tiled_copy_s2t = tcgen05.make_s2t_copy(copy_atom_s2t, tCtSF_compact) + thr_copy_s2t = tiled_copy_s2t.get_slice(0) + + # ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K, STAGE) + tCsSF_compact_s2t_ = thr_copy_s2t.partition_S(tCsSF_compact) + # ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K, STAGE) + tCsSF_compact_s2t = tcgen05.get_s2t_smem_desc_tensor( + tiled_copy_s2t, tCsSF_compact_s2t_ + ) + # ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K) + tCtSF_compact_s2t = thr_copy_s2t.partition_D(tCtSF_compact) + + return tiled_copy_s2t, tCsSF_compact_s2t, tCtSF_compact_s2t + + def epilog_tmem_copy_and_partition( + self, + tidx: cutlass.Int32, + tAcc: cute.Tensor, + gC_mnl: cute.Tensor, + epi_tile: cute.Tile, + use_2cta_instrs: Union[cutlass.Boolean, bool], + ) -> Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]: + """ + Make tiledCopy for tensor memory load, then use it to partition tensor memory (source) and register array (destination). + + :param tidx: The thread index in epilogue warp groups + :type tidx: cutlass.Int32 + :param tAcc: The accumulator tensor to be copied and partitioned + :type tAcc: cute.Tensor + :param gC_mnl: The global tensor C + :type gC_mnl: cute.Tensor + :param epi_tile: The epilogue tiler + :type epi_tile: cute.Tile + :param use_2cta_instrs: Whether use_2cta_instrs is enabled + :type use_2cta_instrs: bool + + :return: A tuple containing (tiled_copy_t2r, tTR_tAcc, tTR_rAcc) where: + - tiled_copy_t2r: The tiled copy operation for tmem to register copy(t2r) + - tTR_tAcc: The partitioned accumulator tensor + - tTR_rAcc: The accumulated tensor in register used to hold t2r results + :rtype: Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor] + """ + # Make tiledCopy for tensor memory load + copy_atom_t2r = sm100_utils.get_tmem_load_op( + self.cta_tile_shape_mnk, + self.c_layout, + self.c_dtype, + self.acc_dtype, + epi_tile, + use_2cta_instrs, + ) + # (EPI_TILE_M, EPI_TILE_N, EPI_M, EPI_N, STAGE) + tAcc_epi = cute.flat_divide( + tAcc[((None, None), 0, 0, None)], + epi_tile, + ) + # (EPI_TILE_M, EPI_TILE_N) + tiled_copy_t2r = tcgen05.make_tmem_copy( + copy_atom_t2r, tAcc_epi[(None, None, 0, 0, 0)] + ) + + thr_copy_t2r = tiled_copy_t2r.get_slice(tidx) + # (T2R, T2R_M, T2R_N, EPI_M, EPI_M, STAGE) + tTR_tAcc = thr_copy_t2r.partition_S(tAcc_epi) + + # (EPI_TILE_M, EPI_TILE_N, EPI_M, EPI_N, RestM, RestN, RestL) + gC_mnl_epi = cute.flat_divide( + gC_mnl[((None, None), 0, 0, None, None, None)], epi_tile + ) + # (T2R, T2R_M, T2R_N, EPI_M, EPI_N, RestM, RestN, RestL) + tTR_gC = thr_copy_t2r.partition_D(gC_mnl_epi) + # (T2R, T2R_M, T2R_N) + tTR_rAcc = cute.make_rmem_tensor( + tTR_gC[(None, None, None, 0, 0, 0, 0, 0)].shape, self.acc_dtype + ) + return tiled_copy_t2r, tTR_tAcc, tTR_rAcc + + def epilog_smem_copy_and_partition( + self, + tiled_copy_t2r: cute.TiledCopy, + tTR_rC: cute.Tensor, + tidx: cutlass.Int32, + sC: cute.Tensor, + ) -> Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]: + """ + Make tiledCopy for shared memory store, then use it to partition register array (source) and shared memory (destination). + + :param tiled_copy_t2r: The tiled copy operation for tmem to register copy(t2r) + :type tiled_copy_t2r: cute.TiledCopy + :param tTR_rC: The partitioned accumulator tensor + :type tTR_rC: cute.Tensor + :param tidx: The thread index in epilogue warp groups + :type tidx: cutlass.Int32 + :param sC: The shared memory tensor to be copied and partitioned + :type sC: cute.Tensor + :type sepi: cute.Tensor + + :return: A tuple containing (tiled_copy_r2s, tRS_rC, tRS_sC) where: + - tiled_copy_r2s: The tiled copy operation for register to smem copy(r2s) + - tRS_rC: The partitioned tensor C (register source) + - tRS_sC: The partitioned tensor C (smem destination) + :rtype: Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor] + """ + copy_atom_r2s = sm100_utils.get_smem_store_op( + self.c_layout, self.c_dtype, self.acc_dtype, tiled_copy_t2r + ) + tiled_copy_r2s = cute.make_tiled_copy_D(copy_atom_r2s, tiled_copy_t2r) + # (R2S, R2S_M, R2S_N, PIPE_D) + thr_copy_r2s = tiled_copy_r2s.get_slice(tidx) + tRS_sC = thr_copy_r2s.partition_D(sC) + # (R2S, R2S_M, R2S_N) + tRS_rC = tiled_copy_r2s.retile(tTR_rC) + return tiled_copy_r2s, tRS_rC, tRS_sC + + def epilog_gmem_copy_and_partition( + self, + tidx: cutlass.Int32, + atom: Union[cute.CopyAtom, cute.TiledCopy], + gC_mnl: cute.Tensor, + epi_tile: cute.Tile, + sC: cute.Tensor, + ) -> Tuple[cute.CopyAtom, cute.Tensor, cute.Tensor]: + """Make tiledCopy for global memory store, then use it to: + partition shared memory (source) and global memory (destination) for TMA store version. + + :param tidx: The thread index in epilogue warp groups + :type tidx: cutlass.Int32 + :param atom: The copy_atom_c to be used for TMA store version, or tiled_copy_t2r for none TMA store version + :type atom: cute.CopyAtom or cute.TiledCopy + :param gC_mnl: The global tensor C + :type gC_mnl: cute.Tensor + :param epi_tile: The epilogue tiler + :type epi_tile: cute.Tile + :param sC: The shared memory tensor to be copied and partitioned + :type sC: cute.Tensor + + :return: A tuple containing (tma_atom_c, bSG_sC, bSG_gC) where: + - tma_atom_c: The TMA copy atom + - bSG_sC: The partitioned shared memory tensor C + - bSG_gC: The partitioned global tensor C + :rtype: Tuple[cute.CopyAtom, cute.Tensor, cute.Tensor] + """ + # (EPI_TILE_M, EPI_TILE_N, EPI_M, EPI_N, RestM, RestN, RestL) + gC_epi = cute.flat_divide( + gC_mnl[((None, None), 0, 0, None, None, None)], epi_tile + ) + + tma_atom_c = atom + sC_for_tma_partition = cute.group_modes(sC, 0, 2) + gC_for_tma_partition = cute.group_modes(gC_epi, 0, 2) + # ((ATOM_V, REST_V), EPI_M, EPI_N) + # ((ATOM_V, REST_V), EPI_M, EPI_N, RestM, RestN, RestL) + bSG_sC, bSG_gC = cpasync.tma_partition( + tma_atom_c, + 0, + cute.make_layout(1), + sC_for_tma_partition, + gC_for_tma_partition, + ) + return tma_atom_c, bSG_sC, bSG_gC + + @staticmethod + def _compute_stages( + tiled_mma: cute.TiledMma, + mma_tiler_mnk: Tuple[int, int, int], + a_dtype: Type[cutlass.Numeric], + b_dtype: Type[cutlass.Numeric], + epi_tile: cute.Tile, + c_dtype: Type[cutlass.Numeric], + c_layout: utils.LayoutEnum, + sf_dtype: Type[cutlass.Numeric], + sf_vec_size: int, + smem_capacity: int, + occupancy: int, + ) -> Tuple[int, int, int]: + """Computes the number of stages for A/B/C operands based on heuristics. + + :param tiled_mma: The tiled MMA object defining the core computation. + :type tiled_mma: cute.TiledMma + :param mma_tiler_mnk: The shape (M, N, K) of the MMA tiler. + :type mma_tiler_mnk: tuple[int, int, int] + :param a_dtype: Data type of operand A. + :type a_dtype: type[cutlass.Numeric] + :param b_dtype: Data type of operand B. + :type b_dtype: type[cutlass.Numeric] + :param epi_tile: The epilogue tile shape. + :type epi_tile: cute.Tile + :param c_dtype: Data type of operand C (output). + :type c_dtype: type[cutlass.Numeric] + :param c_layout: Layout enum of operand C. + :type c_layout: utils.LayoutEnum + :param sf_dtype: Data type of Scale factor. + :type sf_dtype: type[cutlass.Numeric] + :param sf_vec_size: Scale factor vector size. + :type sf_vec_size: int + :param smem_capacity: Total available shared memory capacity in bytes. + :type smem_capacity: int + :param occupancy: Target number of CTAs per SM (occupancy). + :type occupancy: int + + :return: A tuple containing the computed number of stages for: + (ACC stages, A/B operand stages, C stages) + :rtype: tuple[int, int, int] + """ + # ACC stages + num_acc_stage = 1 if mma_tiler_mnk[1] == 256 else 2 + + # Default C stages + num_c_stage = 2 + + # Calculate smem layout and size for one stage of A, B, SFA, SFB and C + a_smem_layout_stage_one = sm100_utils.make_smem_layout_a( + tiled_mma, + mma_tiler_mnk, + a_dtype, + 1, # a tmp 1 stage is provided + ) + b_smem_layout_staged_one = sm100_utils.make_smem_layout_b( + tiled_mma, + mma_tiler_mnk, + b_dtype, + 1, # a tmp 1 stage is provided + ) + sfa_smem_layout_staged_one = blockscaled_utils.make_smem_layout_sfa( + tiled_mma, + mma_tiler_mnk, + sf_vec_size, + 1, # a tmp 1 stage is provided + ) + sfb_smem_layout_staged_one = blockscaled_utils.make_smem_layout_sfb( + tiled_mma, + mma_tiler_mnk, + sf_vec_size, + 1, # a tmp 1 stage is provided + ) + + c_smem_layout_staged_one = sm100_utils.make_smem_layout_epi( + c_dtype, + c_layout, + epi_tile, + 1, + ) + + ab_bytes_per_stage = ( + cute.size_in_bytes(a_dtype, a_smem_layout_stage_one) + + cute.size_in_bytes(b_dtype, b_smem_layout_staged_one) + + cute.size_in_bytes(sf_dtype, sfa_smem_layout_staged_one) + + cute.size_in_bytes(sf_dtype, sfb_smem_layout_staged_one) + ) + mbar_helpers_bytes = 1024 + c_bytes_per_stage = cute.size_in_bytes(c_dtype, c_smem_layout_staged_one) + c_bytes = c_bytes_per_stage * num_c_stage + amax_bytes = 16 + + # Calculate A/B/SFA/SFB stages: + # Start with total smem per CTA (capacity / occupancy) + # Subtract reserved bytes and initial C stages bytes + # Divide remaining by bytes needed per A/B/SFA/SFB stage + num_ab_stage = ( + smem_capacity // occupancy - (mbar_helpers_bytes + c_bytes + amax_bytes) + ) // ab_bytes_per_stage + + # Refine epilogue stages: + # Calculate remaining smem after allocating for A/B/SFA/SFB stages and reserved bytes + # Add remaining unused smem to epilogue + num_c_stage += ( + smem_capacity + - occupancy * ab_bytes_per_stage * num_ab_stage + - occupancy * (mbar_helpers_bytes + c_bytes + amax_bytes) + ) // (occupancy * c_bytes_per_stage) + + return num_acc_stage, num_ab_stage, num_c_stage + + @staticmethod + def _compute_grid( + c: cute.Tensor, + cta_tile_shape_mnk: Tuple[int, int, int], + cluster_shape_mn: Tuple[int, int], + max_active_clusters: cutlass.Constexpr, + ) -> Tuple[utils.PersistentTileSchedulerParams, Tuple[int, int, int]]: + """Use persistent tile scheduler to compute the grid size for the output tensor C. + + :param c: The output tensor C + :type c: cute.Tensor + :param cta_tile_shape_mnk: The shape (M, N, K) of the CTA tile. + :type cta_tile_shape_mnk: tuple[int, int, int] + :param cluster_shape_mn: Shape of each cluster in M, N dimensions. + :type cluster_shape_mn: tuple[int, int] + :param max_active_clusters: Maximum number of active clusters. + :type max_active_clusters: cutlass.Constexpr + + :return: A tuple containing: + - tile_sched_params: Parameters for the persistent tile scheduler. + - grid: Grid shape for kernel launch. + :rtype: Tuple[utils.PersistentTileSchedulerParams, tuple[int, int, int]] + """ + c_shape = cute.slice_(cta_tile_shape_mnk, (None, None, 0)) + gc = cute.zipped_divide(c, tiler=c_shape) + num_ctas_mnl = gc[(0, (None, None, None))].shape + cluster_shape_mnl = (*cluster_shape_mn, 1) + + tile_sched_params = utils.PersistentTileSchedulerParams( + num_ctas_mnl, cluster_shape_mnl + ) + grid = utils.StaticPersistentTileScheduler.get_grid_shape( + tile_sched_params, max_active_clusters + ) + + return tile_sched_params, grid + + @staticmethod + def is_valid_dtypes_and_scale_factor_vec_size( + ab_dtype: Type[cutlass.Numeric], + sf_dtype: Type[cutlass.Numeric], + sf_vec_size: int, + c_dtype: Type[cutlass.Numeric], + ) -> bool: + """ + Check if the dtypes and sf_vec_size are valid combinations + + :param ab_dtype: The data type of the A and B operands + :type ab_dtype: Type[cutlass.Numeric] + :param sf_dtype: The data type of the scale factor + :type sf_dtype: Type[cutlass.Numeric] + :param sf_vec_size: The vector size of the scale factor + :type sf_vec_size: int + :param c_dtype: The data type of the output tensor + :type c_dtype: Type[cutlass.Numeric] + + :return: True if the dtypes and sf_vec_size are valid, False otherwise + :rtype: bool + """ + is_valid = True + + # Check valid ab_dtype + if ab_dtype not in { + cutlass.Float4E2M1FN, + cutlass.Float8E5M2, + cutlass.Float8E4M3FN, + }: + is_valid = False + + # Check valid sf_vec_size + if sf_vec_size not in {16, 32}: + is_valid = False + + # Check valid sf_dtype + if sf_dtype not in {cutlass.Float8E8M0FNU, cutlass.Float8E4M3FN}: + is_valid = False + + # Check valid sf_dtype and sf_vec_size combinations + if sf_dtype == cutlass.Float8E4M3FN and sf_vec_size == 32: + is_valid = False + if ab_dtype in {cutlass.Float8E5M2, cutlass.Float8E4M3FN} and sf_vec_size == 16: + is_valid = False + + # Check valid c_dtype + if c_dtype not in { + cutlass.Float32, + cutlass.Float16, + cutlass.BFloat16, + cutlass.Float8E5M2, + cutlass.Float8E4M3FN, + }: + is_valid = False + + return is_valid + + @staticmethod + def is_valid_layouts( + ab_dtype: Type[cutlass.Numeric], + c_dtype: Type[cutlass.Numeric], + a_major: str, + b_major: str, + c_major: str, + ) -> bool: + """ + Check if layouts and dtypes are valid combinations + + :param ab_dtype: The data type of the A and B operands + :type ab_dtype: Type[cutlass.Numeric] + :param c_dtype: The data type of the output tensor + :type c_dtype: Type[cutlass.Numeric] + :param a_major: The major dimension of the A tensor + :type a_major: str + :param b_major: The major dimension of the B tensor + :type b_major: str + :param c_major: The major dimension of the C tensor + :type c_major: str + + :return: True if the layouts are valid, False otherwise + :rtype: bool + """ + is_valid = True + + if ab_dtype is cutlass.Float4E2M1FN and not (a_major == "k" and b_major == "k"): + is_valid = False + return is_valid + + @staticmethod + def is_valid_mma_tiler_and_cluster_shape( + mma_tiler_mn: Tuple[int, int], + cluster_shape_mn: Tuple[int, int], + ) -> bool: + """ + Check if the mma tiler and cluster shape are valid + + :param mma_tiler_mn: The (M, N) shape of the MMA instruction tiler + :type mma_tiler_mn: Tuple[int, int] + :param cluster_shape_mn: The (ClusterM, ClusterN) shape of the CTA cluster + :type cluster_shape_mn: Tuple[int, int] + + :return: True if the mma tiler and cluster shape are valid, False otherwise + :rtype: bool + """ + is_valid = True + # Skip invalid mma tile shape + if mma_tiler_mn[0] not in [128, 256]: + is_valid = False + if mma_tiler_mn[1] not in [128, 256]: + is_valid = False + # Skip illegal cluster shape + if cluster_shape_mn[0] % (2 if mma_tiler_mn[0] == 256 else 1) != 0: + is_valid = False + # Skip invalid cluster shape + is_power_of_2 = lambda x: x > 0 and (x & (x - 1)) == 0 + if ( + cluster_shape_mn[0] * cluster_shape_mn[1] > 16 + or cluster_shape_mn[0] <= 0 + or cluster_shape_mn[1] <= 0 + # Special cluster shape check for scale factor multicasts. + # Due to limited size of scale factors, we can't multicast among more than 4 CTAs. + or cluster_shape_mn[0] > 4 + or cluster_shape_mn[1] > 4 + or not is_power_of_2(cluster_shape_mn[0]) + or not is_power_of_2(cluster_shape_mn[1]) + ): + is_valid = False + return is_valid + + @staticmethod + def is_valid_tensor_alignment( + m: int, + n: int, + k: int, + l: int, + ab_dtype: Type[cutlass.Numeric], + c_dtype: Type[cutlass.Numeric], + a_major: str, + b_major: str, + c_major: str, + ) -> bool: + """ + Check if the tensor alignment is valid + + :param m: The number of rows in the A tensor + :type m: int + :param n: The number of columns in the B tensor + :type n: int + :param k: The number of columns in the A tensor + :type k: int + :param l: The number of columns in the C tensor + :type l: int + :param ab_dtype: The data type of the A and B operands + :type ab_dtype: Type[cutlass.Numeric] + :param c_dtype: The data type of the output tensor + :type c_dtype: Type[cutlass.Numeric] + :param a_major: The major axis of the A tensor + :type a_major: str + :param b_major: The major axis of the B tensor + :type b_major: str + :param c_major: The major axis of the C tensor + :type c_major: str + + :return: True if the problem shape is valid, False otherwise + :rtype: bool + """ + is_valid = True + + def check_contigous_16B_alignment(dtype, is_mode0_major, tensor_shape): + major_mode_idx = 0 if is_mode0_major else 1 + num_major_elements = tensor_shape[major_mode_idx] + num_contiguous_elements = 16 * 8 // dtype.width + return num_major_elements % num_contiguous_elements == 0 + + if ( + not check_contigous_16B_alignment(ab_dtype, a_major == "m", (m, k, l)) + or not check_contigous_16B_alignment(ab_dtype, b_major == "n", (n, k, l)) + or not check_contigous_16B_alignment(c_dtype, c_major == "m", (m, n, l)) + ): + is_valid = False + return is_valid + + @staticmethod + def can_implement( + ab_dtype: Type[cutlass.Numeric], + sf_dtype: Type[cutlass.Numeric], + sf_vec_size: int, + c_dtype: Type[cutlass.Numeric], + mma_tiler_mn: Tuple[int, int], + cluster_shape_mn: Tuple[int, int], + m: int, + n: int, + k: int, + l: int, + a_major: str, + b_major: str, + c_major: str, + ) -> bool: + """ + Check if the gemm can be implemented + + :param ab_dtype: The data type of the A and B operands + :type ab_dtype: Type[cutlass.Numeric] + :param sf_dtype: The data type of the scale factor tensor + :type sf_dtype: Type[cutlass.Numeric] + :param sf_vec_size: The vector size + :type sf_vec_size: int + :param c_dtype: The data type of the output tensor + :type c_dtype: Type[cutlass.Numeric] + :param mma_tiler_mn: The (M, N) shape of the MMA instruction tiler + :type mma_tiler_mn: Tuple[int, int] + :param cluster_shape_mn: The (ClusterM, ClusterN) shape of the CTA cluster + :type cluster_shape_mn: Tuple[int, int] + :param m: The number of rows in the A tensor + :type m: int + :param n: The number of columns in the B tensor + :type n: int + :param k: The number of columns in the A tensor + :type k: int + :param l: The number of columns in the C tensor + :type l: int + :param a_major: The major axis of the A tensor + :type a_major: str + :param b_major: The major axis of the B tensor + :type b_major: str + :param c_major: The major axis of the C tensor + :type c_major: str + + :return: True if the gemm can be implemented, False otherwise + :rtype: bool + """ + can_implement = True + # Skip unsupported types + if not Sm100BlockScaledPersistentDenseGemmKernel.is_valid_dtypes_and_scale_factor_vec_size( + ab_dtype, sf_dtype, sf_vec_size, c_dtype + ): + can_implement = False + # Skip unsupported layouts + if not Sm100BlockScaledPersistentDenseGemmKernel.is_valid_layouts( + ab_dtype, c_dtype, a_major, b_major, c_major + ): + can_implement = False + # Skip invalid mma tile shape and cluster shape + if not Sm100BlockScaledPersistentDenseGemmKernel.is_valid_mma_tiler_and_cluster_shape( + mma_tiler_mn, cluster_shape_mn + ): + can_implement = False + # Skip illegal problem shape for load/store alignment + if not Sm100BlockScaledPersistentDenseGemmKernel.is_valid_tensor_alignment( + m, n, k, l, ab_dtype, c_dtype, a_major, b_major, c_major + ): + can_implement = False + return can_implement + + +@cute.jit +def cvt_sf_MKL_to_M32x4xrm_K4xrk_L( + sf_ref_tensor: cute.Tensor, + sf_mma_tensor: cute.Tensor, +): + """Convert scale factor tensor from MKL layout to mma specification M(32x4xrest_m)xK(4xrest_k)xL layout""" + # sf_mma_tensor has flatten shape (32, 4, rest_m, 4, rest_k, l) + # group to ((32, 4, rest_m), (4, rest_k), l) + sf_mma_tensor = cute.group_modes(sf_mma_tensor, 0, 3) + sf_mma_tensor = cute.group_modes(sf_mma_tensor, 1, 3) + for i in cutlass.range(cute.size(sf_ref_tensor)): + mkl_coord = sf_ref_tensor.layout.get_hier_coord(i) + sf_mma_tensor[mkl_coord] = sf_ref_tensor[mkl_coord] + + +def compute_reference_amax(output_tensor) -> float: + import torch + + """ + Compute reference amax value on CPU. + + Args: + output_tensor: torch.Tensor, GEMM output result (CPU tensor) + + Returns: + float: reference amax value + """ + # Ensure FP32 for computation + if output_tensor.dtype != torch.float32: + output_fp32 = output_tensor.float() + else: + output_fp32 = output_tensor + + # Compute absolute maximum value + reference_amax = torch.amax(torch.abs(output_fp32)) + + return reference_amax.item() + + +def run( + mnkl: Tuple[int, int, int, int], + ab_dtype: Type[cutlass.Numeric], + sf_dtype: Type[cutlass.Numeric], + sf_vec_size: int, + c_dtype: Type[cutlass.Numeric], + a_major: str, + b_major: str, + c_major: str, + mma_tiler_mn: Tuple[int, int], + cluster_shape_mn: Tuple[int, int], + tolerance: float = 1e-01, + warmup_iterations: int = 0, + iterations: int = 1, + skip_ref_check: bool = False, + use_cold_l2: bool = False, + **kwargs, +): + """Execute a persistent batched dense blockscaled GEMM operation on Blackwell architecture with performance benchmarking. + + This function prepares input tensors, configures and launches the persistent GEMM kernel, + optionally performs reference validation, and benchmarks the execution performance. + + :param mnkl: Problem size (M, N, K, L) + :type mnkl: Tuple[int, int, int, int] + :param ab_dtype: Data type for input tensors A and B + :type ab_dtype: Type[cutlass.Numeric] + :param sf_dtype: Data type for scale factor tensor + :type sf_dtype: Type[cutlass.Numeric] + :param sf_vec_size: Vector size for scale factor tensor + :type sf_vec_size: int + :param c_dtype: Data type for output tensor C + :type c_dtype: Type[cutlass.Numeric] + :param a_major/b_major/c_major: Memory layout of tensor A/B/C + :type a_major/b_major/c_major: str + :param mma_tiler_mn: MMA tiling size. + :type mma_tiler_mn: Tuple[int, int] + :param cluster_shape_mn: Cluster shape. + :type cluster_shape_mn: Tuple[int, int] + :param tolerance: Tolerance value for reference validation comparison, defaults to 1e-01 + :type tolerance: float, optional + :param warmup_iterations: Number of warmup iterations before benchmarking, defaults to 0 + :type warmup_iterations: int, optional + :param iterations: Number of benchmark iterations to run, defaults to 1 + :type iterations: int, optional + :param skip_ref_check: Whether to skip reference result validation, defaults to False + :type skip_ref_check: bool, optional + :param use_cold_l2: Whether to use circular buffer strategy to ensure cold L2 cache, defaults to False + :type use_cold_l2: bool, optional + :raises RuntimeError: If CUDA GPU is not available + :raises ValueError: If the configuration is invalid or unsupported by the kernel + :return: Execution time of the GEMM kernel + :rtype: float + """ + print("Running Sm100 Persistent Dense BlockScaled GEMM test with:") + print(f"mnkl: {mnkl}") + print(f"AB dtype: {ab_dtype}, SF dtype: {sf_dtype}, SF Vec size: {sf_vec_size}") + print(f"C dtype: {c_dtype}") + print(f"Matrix majors - A: {a_major}, B: {b_major}, C: {c_major}") + print(f"Mma Tiler (M, N): {mma_tiler_mn}, Cluster Shape (M, N): {cluster_shape_mn}") + print(f"Tolerance: {tolerance}") + print(f"Warmup iterations: {warmup_iterations}") + print(f"Iterations: {iterations}") + print(f"Skip reference checking: {skip_ref_check}") + print(f"Use cold L2: {'True' if use_cold_l2 else 'False'}") + import torch + import cutlass.torch as cutlass_torch + + # Unpack parameters + m, n, k, l = mnkl + + # Skip unsupported testcase + if not Sm100BlockScaledPersistentDenseGemmKernel.can_implement( + ab_dtype, + sf_dtype, + sf_vec_size, + c_dtype, + mma_tiler_mn, + cluster_shape_mn, + m, + n, + k, + l, + a_major, + b_major, + c_major, + ): + raise TypeError( + f"Unsupported testcase {ab_dtype}, {sf_dtype}, {sf_vec_size}, {c_dtype}, {mma_tiler_mn}, {cluster_shape_mn}, {m}, {n}, {k}, {l}, {a_major}, {b_major}, {c_major}" + ) + + if not torch.cuda.is_available(): + raise RuntimeError("GPU is required to run this example!") + + torch.manual_seed(1111) + + # Create tensor A/B/C + a_ref = cutlass_torch.matrix(l, m, k, a_major == "m", cutlass.Float32) + b_ref = cutlass_torch.matrix(l, n, k, b_major == "n", cutlass.Float32) + c_ref = cutlass_torch.matrix(l, m, n, c_major == "m", cutlass.Float32) + + a_tensor, a_torch = cutlass_torch.cute_tensor_like( + a_ref, ab_dtype, is_dynamic_layout=True, assumed_align=16 + ) + b_tensor, b_torch = cutlass_torch.cute_tensor_like( + b_ref, ab_dtype, is_dynamic_layout=True, assumed_align=16 + ) + c_tensor, c_torch = cutlass_torch.cute_tensor_like( + c_ref, c_dtype, is_dynamic_layout=True, assumed_align=16 + ) + + # Create amax tensor (single FP32 value, initialized to -inf) + amax_ref = cutlass_torch.matrix( + 1, + 1, + 1, + False, + cutlass.Float32, + init_type=cutlass_torch.TensorInitType.SCALAR, + init_config=cutlass_torch.ScalarInitConfig(-float("inf")), + ) + amax_tensor, amax_torch = cutlass_torch.cute_tensor_like( + amax_ref, cutlass.Float32, is_dynamic_layout=True, assumed_align=16 + ) + + # Mark tensor with element divisibility for 16B alignment + a_tensor.mark_compact_shape_dynamic( + mode=1 if a_major == "k" else 0, + stride_order=(2, 0, 1) if a_major == "k" else (2, 1, 0), + divisibility=32 if ab_dtype == cutlass.Float4E2M1FN else 16, + ) + b_tensor.mark_compact_shape_dynamic( + mode=1 if b_major == "k" else 0, + stride_order=(2, 0, 1) if b_major == "k" else (2, 1, 0), + divisibility=32 if ab_dtype == cutlass.Float4E2M1FN else 16, + ) + c_tensor.mark_compact_shape_dynamic( + mode=1 if c_major == "n" else 0, + stride_order=(2, 0, 1) if c_major == "n" else (2, 1, 0), + divisibility=32 if ab_dtype == cutlass.Float4E2M1FN else 16, + ) + + # Create scale factor tensor SFA/SFB + def create_scale_factor_tensor(l, mn, k, sf_vec_size, dtype): + def ceil_div(a, b): + return (a + b - 1) // b + + sf_k = ceil_div(k, sf_vec_size) + ref_shape = (l, mn, sf_k) + + atom_m = (32, 4) + atom_k = 4 + mma_shape = ( + l, + ceil_div(mn, atom_m[0] * atom_m[1]), + ceil_div(sf_k, atom_k), + atom_m[0], + atom_m[1], + atom_k, + ) + + ref_permute_order = (1, 2, 0) + mma_permute_order = (3, 4, 1, 5, 2, 0) + + # Create f32 ref torch tensor (cpu) + ref_f32_torch_tensor_cpu = cutlass_torch.create_and_permute_torch_tensor( + ref_shape, + torch.float32, + permute_order=ref_permute_order, + init_type=cutlass_torch.TensorInitType.RANDOM, + init_config=cutlass_torch.RandomInitConfig( + min_val=1, + max_val=3, + ), + ) + + # Create f32 cute torch tensor (cpu) + cute_f32_torch_tensor_cpu = cutlass_torch.create_and_permute_torch_tensor( + mma_shape, + torch.float32, + permute_order=mma_permute_order, + init_type=cutlass_torch.TensorInitType.RANDOM, + init_config=cutlass_torch.RandomInitConfig( + min_val=0, + max_val=1, + ), + ) + + # convert ref f32 tensor to cute f32 tensor + cvt_sf_MKL_to_M32x4xrm_K4xrk_L( + from_dlpack(ref_f32_torch_tensor_cpu), + from_dlpack(cute_f32_torch_tensor_cpu), + ) + cute_f32_torch_tensor = cute_f32_torch_tensor_cpu.cuda() + + # reshape makes memory contiguous + ref_f32_torch_tensor_cpu = ( + ref_f32_torch_tensor_cpu.permute(2, 0, 1) + .unsqueeze(-1) + .expand(l, mn, sf_k, sf_vec_size) + .reshape(l, mn, sf_k * sf_vec_size) + .permute(*ref_permute_order) + ) + # prune to mkl for reference check. + ref_f32_torch_tensor_cpu = ref_f32_torch_tensor_cpu[:, :k, :] + + # Create dtype cute torch tensor (cpu) + cute_tensor, cute_torch_tensor = cutlass_torch.cute_tensor_like( + cute_f32_torch_tensor_cpu, + dtype, + is_dynamic_layout=True, + assumed_align=16, + ) + + # Convert f32 cute tensor to dtype cute tensor + cute_tensor = cutlass_torch.convert_cute_tensor( + cute_f32_torch_tensor, + cute_tensor, + dtype, + is_dynamic_layout=True, + ) + return ref_f32_torch_tensor_cpu, cute_tensor, cute_torch_tensor + + sfa_ref, sfa_tensor, sfa_torch = create_scale_factor_tensor( + l, m, k, sf_vec_size, sf_dtype + ) + sfb_ref, sfb_tensor, sfb_torch = create_scale_factor_tensor( + l, n, k, sf_vec_size, sf_dtype + ) + + # Configure gemm kernel + gemm = Sm100BlockScaledPersistentDenseGemmKernel( + sf_vec_size, + mma_tiler_mn, + cluster_shape_mn, + ) + + # Compute max active clusters on current device + hardware_info = cutlass.utils.HardwareInfo() + max_active_clusters = hardware_info.get_max_active_clusters( + cluster_shape_mn[0] * cluster_shape_mn[1] + ) + + # Initialize Stream + current_stream = cutlass_torch.default_stream() + + # Compile gemm kernel + compiled_gemm = cute.compile( + gemm, + a_tensor, + b_tensor, + sfa_tensor, + sfb_tensor, + c_tensor, + amax_tensor, + max_active_clusters, + current_stream, + options=f"--opt-level 2", + ) + + # Compute reference result + if not skip_ref_check: + # Execute kernel once for reference checking + compiled_gemm( + a_tensor, + b_tensor, + sfa_tensor, + sfb_tensor, + c_tensor, + amax_tensor, + current_stream, + ) + print("Verifying results...") + res_a = torch.einsum("mkl,mkl->mkl", a_ref, sfa_ref) + res_b = torch.einsum("nkl,nkl->nkl", b_ref, sfb_ref) + ref = torch.einsum("mkl,nkl->mnl", res_a, res_b) + + # Save original Float32 ref for amax computation (before quantization) + ref_for_amax = ref.clone() + + # Convert c back to f32 for comparison. + c_ref_device = c_ref.cuda() + cute.testing.convert( + c_tensor, + from_dlpack(c_ref_device, assumed_align=16).mark_layout_dynamic( + leading_dim=(1 if c_major == "n" else 0) + ), + ) + c_ref = c_ref_device.cpu() + + if c_dtype in (cutlass.Float32, cutlass.Float16, cutlass.BFloat16): + torch.testing.assert_close(c_ref, ref, atol=tolerance, rtol=1e-02) + elif c_dtype in (cutlass.Float8E5M2, cutlass.Float8E4M3FN): + # Convert ref : f32 -> f8 -> f32 + ref_f8_ = torch.empty(*(l, m, n), dtype=torch.uint8, device="cuda").permute( + 1, 2, 0 + ) + ref_f8 = from_dlpack(ref_f8_, assumed_align=16).mark_layout_dynamic( + leading_dim=1 + ) + ref_f8.element_type = c_dtype + ref_device = ref.permute(2, 0, 1).contiguous().permute(1, 2, 0).cuda() + ref_tensor = from_dlpack(ref_device, assumed_align=16).mark_layout_dynamic( + leading_dim=1 + ) + cute.testing.convert(ref_tensor, ref_f8) + cute.testing.convert(ref_f8, ref_tensor) + ref = ref_device.cpu() + torch.testing.assert_close(c_ref, ref, atol=tolerance, rtol=1e-02) + # Verify amax result + device_amax = amax_torch.cpu().squeeze() # Remove dimensions to make it scalar + # Use original Float32 ref (before quantization) for amax computation + reference_amax = torch.tensor(compute_reference_amax(ref_for_amax)) + + # AMAX validation using same approach as GEMM result validation + torch.testing.assert_close( + device_amax, reference_amax, atol=tolerance, rtol=1e-02 + ) + + def generate_tensors(): + a_tensor, _ = cutlass_torch.cute_tensor_like( + a_ref, ab_dtype, is_dynamic_layout=True, assumed_align=16 + ) + b_tensor, _ = cutlass_torch.cute_tensor_like( + b_ref, ab_dtype, is_dynamic_layout=True, assumed_align=16 + ) + c_tensor, _ = cutlass_torch.cute_tensor_like( + c_ref, c_dtype, is_dynamic_layout=True, assumed_align=16 + ) + + # Mark tensor to be byte aligned + a_tensor.mark_compact_shape_dynamic( + mode=1 if a_major == "k" else 0, + stride_order=(2, 0, 1) if a_major == "k" else (2, 1, 0), + divisibility=2 if ab_dtype == cutlass.Float4E2M1FN else 1, + ) + b_tensor.mark_compact_shape_dynamic( + mode=1 if b_major == "k" else 0, + stride_order=(2, 0, 1) if b_major == "k" else (2, 1, 0), + divisibility=2 if ab_dtype == cutlass.Float4E2M1FN else 1, + ) + c_tensor.mark_compact_shape_dynamic( + mode=1 if c_major == "n" else 0, + stride_order=(2, 0, 1) if c_major == "n" else (2, 1, 0), + divisibility=2 if c_dtype == cutlass.Float4E2M1FN else 1, + ) + + _, sfa_tensor, _ = create_scale_factor_tensor(l, m, k, sf_vec_size, sf_dtype) + _, sfb_tensor, _ = create_scale_factor_tensor(l, n, k, sf_vec_size, sf_dtype) + + # Create amax tensor for benchmarking (reset to -inf for each iteration) + amax_ref_bench = cutlass_torch.matrix( + 1, + 1, + 1, + False, + cutlass.Float32, + init_type=cutlass_torch.TensorInitType.SCALAR, + init_config=cutlass_torch.ScalarInitConfig(-float("inf")), + ) + amax_tensor_bench, _ = cutlass_torch.cute_tensor_like( + amax_ref_bench, cutlass.Float32, is_dynamic_layout=True, assumed_align=16 + ) + + return cute.testing.JitArguments( + a_tensor, + b_tensor, + sfa_tensor, + sfb_tensor, + c_tensor, + amax_tensor_bench, + current_stream, + ) + + workspace_count = 1 + if use_cold_l2: + one_workspace_bytes = ( + a_torch.numel() * a_torch.element_size() + + b_torch.numel() * b_torch.element_size() + + sfa_torch.numel() * sfa_torch.element_size() + + sfb_torch.numel() * sfb_torch.element_size() + + c_torch.numel() * c_torch.element_size() + ) + workspace_count = cute.testing.get_workspace_count( + one_workspace_bytes, warmup_iterations, iterations + ) + + exec_time = cute.testing.benchmark( + compiled_gemm, + workspace_generator=generate_tensors, + workspace_count=workspace_count, + stream=current_stream, + warmup_iterations=warmup_iterations, + iterations=iterations, + ) + + return exec_time # Return execution time in microseconds + + +if __name__ == "__main__": + + def parse_comma_separated_ints(s: str) -> Tuple[int, ...]: + try: + return tuple(int(x.strip()) for x in s.split(",")) + except ValueError: + raise argparse.ArgumentTypeError( + "Invalid format. Expected comma-separated integers." + ) + + parser = argparse.ArgumentParser( + description="Example of Sm100 Dense Persistent BlockScaled GEMM." + ) + + parser.add_argument( + "--mnkl", + type=parse_comma_separated_ints, + default=(512, 256, 256, 1), + help="mnkl dimensions (comma-separated)", + ) + parser.add_argument( + "--mma_tiler_mn", + type=parse_comma_separated_ints, + default=(128, 128), + help="Mma tile shape (comma-separated)", + ) + parser.add_argument( + "--cluster_shape_mn", + type=parse_comma_separated_ints, + default=(1, 1), + help="Cluster shape (comma-separated)", + ) + parser.add_argument("--ab_dtype", type=cutlass.dtype, default=cutlass.Float4E2M1FN) + parser.add_argument("--sf_dtype", type=cutlass.dtype, default=cutlass.Float8E8M0FNU) + parser.add_argument("--sf_vec_size", type=int, default=16) + parser.add_argument("--c_dtype", type=cutlass.dtype, default=cutlass.Float16) + parser.add_argument("--a_major", choices=["k", "m"], type=str, default="k") + parser.add_argument("--b_major", choices=["k", "n"], type=str, default="k") + parser.add_argument("--c_major", choices=["n", "m"], type=str, default="n") + parser.add_argument( + "--tolerance", type=float, default=1e-01, help="Tolerance for validation" + ) + parser.add_argument( + "--warmup_iterations", type=int, default=0, help="Warmup iterations" + ) + parser.add_argument( + "--iterations", + type=int, + default=1, + help="Number of iterations to run the kernel", + ) + parser.add_argument( + "--skip_ref_check", action="store_true", help="Skip reference checking" + ) + parser.add_argument( + "--use_cold_l2", + action="store_true", + default=False, + help="Use circular buffer tensor sets to ensure L2 cold cache", + ) + + args = parser.parse_args() + + if len(args.mnkl) != 4: + parser.error("--mnkl must contain exactly 4 values") + + if len(args.mma_tiler_mn) != 2: + parser.error("--mma_tiler_mn must contain exactly 2 values") + + if len(args.cluster_shape_mn) != 2: + parser.error("--cluster_shape_mn must contain exactly 2 values") + + run( + args.mnkl, + args.ab_dtype, + args.sf_dtype, + args.sf_vec_size, + args.c_dtype, + args.a_major, + args.b_major, + args.c_major, + args.mma_tiler_mn, + args.cluster_shape_mn, + args.tolerance, + args.warmup_iterations, + args.iterations, + args.skip_ref_check, + args.use_cold_l2, + ) + print("PASS") diff --git a/examples/python/CuTeDSL/blackwell/dense_gemm.py b/examples/python/CuTeDSL/blackwell/dense_gemm.py index 1e347bd1..94495fb1 100644 --- a/examples/python/CuTeDSL/blackwell/dense_gemm.py +++ b/examples/python/CuTeDSL/blackwell/dense_gemm.py @@ -30,7 +30,6 @@ import argparse from typing import Optional, Type, Tuple, Union import cuda.bindings.driver as cuda -import torch import cutlass import cutlass.cute as cute @@ -38,7 +37,6 @@ import cutlass.utils as utils import cutlass.pipeline as pipeline from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait from cutlass.cute.nvgpu import cpasync, tcgen05 -import cutlass.torch as cutlass_torch import cutlass.utils.blackwell_helpers as sm100_utils import cutlass.cute.testing as testing @@ -215,7 +213,6 @@ class DenseGemmKernel: self.occupancy = 1 self.threads_per_cta = 128 - self.smem_capacity = utils.get_smem_capacity_in_bytes() def _setup_attributes(self): """Set up configurations that are dependent on GEMM inputs @@ -278,6 +275,8 @@ class DenseGemmKernel: else: self.epi_tile = self.cta_tile_shape_mnk[:2] + self.smem_capacity = utils.get_smem_capacity_in_bytes() + # Setup A/B/C stage count in shared memory self.num_acc_stage, self.num_ab_stage, self.num_c_stage = self._compute_stages( tiled_mma, @@ -1030,7 +1029,10 @@ class DenseGemmKernel: c_buffer = subtile_idx % self.num_c_stage cute.copy(tiled_copy_r2s, tRS_rC, tRS_sC[(None, None, None, c_buffer)]) # Fence and barrier to make sure shared memory store is visible to TMA store - cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) pipeline.sync(barrier_id=1) # TMA store C to global memory @@ -1261,7 +1263,7 @@ class DenseGemmKernel: """ acc_shape = tiled_mma.partition_shape_C(mma_tiler[:2]) tCtAcc_fake = tiled_mma.make_fragment_C(acc_shape) - return sm100_utils.get_num_tmem_alloc_cols(tCtAcc_fake) + return utils.get_num_tmem_alloc_cols(tCtAcc_fake) def is_valid_dtypes( self, ab_dtype: Type[cutlass.Numeric], c_dtype: Type[cutlass.Numeric] @@ -1494,6 +1496,9 @@ class DenseGemmKernel: def create_tensors(l, m, n, k, a_major, b_major, c_major, ab_dtype, c_dtype): + import torch + import cutlass.torch as cutlass_torch + torch.manual_seed(1111) a_torch_cpu = cutlass_torch.matrix(l, m, k, a_major == "m", ab_dtype) @@ -1522,6 +1527,9 @@ def create_tensors(l, m, n, k, a_major, b_major, c_major, ab_dtype, c_dtype): def compare(a_torch_cpu, b_torch_cpu, c_torch_gpu, c_dtype, tolerance): + import torch + import cutlass.torch as cutlass_torch + # Copy gpu result back kernel_result = c_torch_gpu.cpu() @@ -1615,6 +1623,7 @@ def run( print(f"Iterations: {iterations}") print(f"Skip reference checking: {skip_ref_check}") print(f"Use cold L2: {'True' if use_cold_l2 else 'False'}") + import torch # Unpack parameters m, n, k, l = mnkl @@ -1654,6 +1663,8 @@ def run( compare(a_torch_cpu, b_torch_cpu, c_torch_gpu, c_dtype, tolerance) def generate_tensors(): + import cutlass.torch as cutlass_torch + a_tensor, _ = cutlass_torch.cute_tensor_like( a_torch_cpu, ab_dtype, is_dynamic_layout=True, assumed_align=16 ) diff --git a/examples/python/CuTeDSL/blackwell/dense_gemm_alpha_beta_persistent.py b/examples/python/CuTeDSL/blackwell/dense_gemm_alpha_beta_persistent.py index 2b7d7d84..d8881325 100644 --- a/examples/python/CuTeDSL/blackwell/dense_gemm_alpha_beta_persistent.py +++ b/examples/python/CuTeDSL/blackwell/dense_gemm_alpha_beta_persistent.py @@ -29,13 +29,11 @@ import argparse from typing import Optional, Tuple, Type, Union -import torch import cuda.bindings.driver as cuda import cutlass import cutlass.cute as cute import cutlass.cute.testing as testing -import cutlass.torch as cutlass_torch import cutlass.utils as utils import cutlass.pipeline as pipeline from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait @@ -1176,7 +1174,10 @@ class SM100PersistentDenseGemmAlphaBetaKernel: tSR_sC[(None, None, None, c_pipeline_consumer_state.index)], tSR_rC, ) - cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) c_pipeline.consumer_release(c_pipeline_consumer_state) # Advance pipeline states @@ -1203,7 +1204,10 @@ class SM100PersistentDenseGemmAlphaBetaKernel: tiled_copy_r2s, tRS_rD, tRS_sD[(None, None, None, d_buffer)] ) # Fence and barrier to make sure shared memory store is visible to TMA store - cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) epilog_sync_barrier.arrive_and_wait() # @@ -1875,6 +1879,9 @@ class SM100PersistentDenseGemmAlphaBetaKernel: def create_tensors(l, m, n, k, a_major, b_major, cd_major, ab_dtype, c_dtype, d_dtype): + import torch + import cutlass.torch as cutlass_torch + torch.manual_seed(1111) a_torch_cpu = cutlass_torch.matrix(l, m, k, a_major == "m", ab_dtype) @@ -1991,6 +1998,9 @@ def compare( beta, tolerance, ): + import torch + import cutlass.torch as cutlass_torch + # Copy gpu result back kernel_result = d_torch_gpu.cpu() @@ -2057,6 +2067,8 @@ def run_dense_gemm( # Unpack parameters m, n, k, l = mnkl + import torch + if not torch.cuda.is_available(): raise RuntimeError("GPU is required to run this example!") diff --git a/examples/python/CuTeDSL/blackwell/dense_gemm_persistent.py b/examples/python/CuTeDSL/blackwell/dense_gemm_persistent.py index e54b1412..d08ea62b 100644 --- a/examples/python/CuTeDSL/blackwell/dense_gemm_persistent.py +++ b/examples/python/CuTeDSL/blackwell/dense_gemm_persistent.py @@ -28,14 +28,14 @@ import argparse from typing import Optional, Tuple, Type, Union - +from functools import lru_cache import cuda.bindings.driver as cuda import cutlass import cutlass.cute as cute import cutlass.cute.testing as testing -from cutlass.cute.runtime import from_dlpack import cutlass.utils as utils +from cutlass.utils import is_fp8_dtype, create_cute_tensor_for_fp8 import cutlass.pipeline as pipeline from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait from cutlass.cute.nvgpu import cpasync, tcgen05 @@ -281,6 +281,7 @@ class PersistentDenseGemmKernel: self.mma_tiler_mn = mma_tiler_mn self.mma_tiler = (*mma_tiler_mn, 1) self.use_tma_store = use_tma_store + self.arch = "sm_100" self.cta_group = ( tcgen05.CtaGroup.TWO if use_2cta_instrs else tcgen05.CtaGroup.ONE @@ -288,17 +289,26 @@ class PersistentDenseGemmKernel: self.occupancy = 1 # Set specialized warp ids - self.epilog_warp_id = (0, 1, 2, 3) + self.epilogue_warp_id = (0, 1, 2, 3) self.mma_warp_id = 4 self.tma_warp_id = 5 self.threads_per_cta = 32 * len( - (self.mma_warp_id, self.tma_warp_id, *self.epilog_warp_id) + (self.mma_warp_id, self.tma_warp_id, *self.epilogue_warp_id) ) # Set barrier id for cta sync, epilogue sync and tmem ptr sync self.epilog_sync_bar_id = 1 self.tmem_alloc_sync_bar_id = 2 self.tmem_dealloc_sync_bar_id = 3 - self.smem_capacity = utils.get_smem_capacity_in_bytes("sm_100") + + def _create_tiled_mma(self): + return utils.sm100.make_trivial_tiled_mma( + self.a_dtype, + self.a_major_mode, + self.b_major_mode, + self.acc_dtype, + self.cta_group, + self.mma_tiler[:2], + ) def _setup_attributes(self): """Set up configurations that are dependent on GEMM inputs @@ -315,14 +325,7 @@ class PersistentDenseGemmKernel: - Computing tensor memory allocation columns """ # Configure tiled mma - tiled_mma = utils.sm100.make_trivial_tiled_mma( - self.a_dtype, - self.a_major_mode, - self.b_major_mode, - self.acc_dtype, - self.cta_group, - self.mma_tiler[:2], - ) + tiled_mma = self._create_tiled_mma() # Compute mma/cluster/tile shapes mma_inst_shape_k = cute.size(tiled_mma.shape_mnk, mode=[2]) @@ -367,6 +370,8 @@ class PersistentDenseGemmKernel: self.c_dtype, self.c_layout, self.epi_tile, 1 ) + self.smem_capacity = utils.get_smem_capacity_in_bytes() + # Setup A/B/C stage count in shared memory and ACC stage count in tensor memory self.num_acc_stage, self.num_ab_stage, self.num_c_stage = _compute_stages( tiled_mma, @@ -396,7 +401,7 @@ class PersistentDenseGemmKernel: # Compute the number of tensor memory allocation columns self.num_tmem_alloc_cols = self._compute_num_tmem_alloc_cols( - tiled_mma, self.mma_tiler, self.num_acc_stage + tiled_mma, self.mma_tiler, self.num_acc_stage, self.arch ) @cute.jit @@ -443,17 +448,11 @@ class PersistentDenseGemmKernel: if cutlass.const_expr(self.a_dtype != self.b_dtype): raise TypeError(f"Type must match: {self.a_dtype} != {self.b_dtype}") + tiled_mma = self._create_tiled_mma() + # Setup attributes that dependent on gemm inputs self._setup_attributes() - tiled_mma = utils.sm100.make_trivial_tiled_mma( - self.a_dtype, - self.a_major_mode, - self.b_major_mode, - self.acc_dtype, - self.cta_group, - self.mma_tiler[:2], - ) atom_thr_size = cute.size(tiled_mma.thr_id.shape) # Setup TMA load for A @@ -618,7 +617,7 @@ class PersistentDenseGemmKernel: # Initialize acc_pipeline (barrier) and states acc_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) - num_acc_consumer_threads = len(self.epilog_warp_id) * ( + num_acc_consumer_threads = len(self.epilogue_warp_id) * ( 2 if use_2cta_instrs else 1 ) acc_pipeline_consumer_group = pipeline.CooperativeGroup( @@ -635,19 +634,19 @@ class PersistentDenseGemmKernel: tmem_alloc_barrier = pipeline.NamedBarrier( barrier_id=self.tmem_alloc_sync_bar_id, - num_threads=32 * len((self.mma_warp_id, *self.epilog_warp_id)), + num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id)), ) tmem_dealloc_barrier = None if cutlass.const_expr(not self.use_tma_store): tmem_dealloc_barrier = pipeline.NamedBarrier( barrier_id=self.tmem_dealloc_sync_bar_id, - num_threads=32 * len(self.epilog_warp_id), + num_threads=32 * len(self.epilogue_warp_id), ) # Tensor memory dealloc barrier init tmem = utils.TmemAllocator( storage.tmem_holding_buf, barrier_for_retrieve=tmem_alloc_barrier, - allocator_warp_id=self.epilog_warp_id[0], + allocator_warp_id=self.epilogue_warp_id[0], is_two_cta=use_2cta_instrs, two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar_ptr, ) @@ -763,6 +762,16 @@ class PersistentDenseGemmKernel: # pipeline_init_wait(cluster_shape_mn=cluster_layout_vmnk) + # + # Construct the scheduler + # + tile_sched = utils.StaticPersistentTileScheduler.create( + tile_sched_params, + cute.arch.block_idx(), + cute.arch.grid_dim(), + ) + work_tile = tile_sched.initial_work_tile_info() + # # Specialized TMA load warp # @@ -771,10 +780,6 @@ class PersistentDenseGemmKernel: # # Persistent tile scheduling loop # - tile_sched = utils.StaticPersistentTileScheduler.create( - tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim() - ) - work_tile = tile_sched.initial_work_tile_info() while work_tile.is_valid_tile: # Get tile coord from tile scheduler @@ -855,10 +860,6 @@ class PersistentDenseGemmKernel: # # Persistent tile scheduling loop # - tile_sched = utils.StaticPersistentTileScheduler.create( - tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim() - ) - work_tile = tile_sched.initial_work_tile_info() acc_producer_state = pipeline.make_pipeline_state( pipeline.PipelineUserType.Producer, self.num_acc_stage @@ -973,37 +974,70 @@ class PersistentDenseGemmKernel: # # Persistent tile scheduling loop for epilogue # - tile_sched = utils.StaticPersistentTileScheduler.create( - tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim() + acc_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_acc_stage ) if cutlass.const_expr(self.use_tma_store): assert tma_atom_c is not None and sC is not None - self.epilogue_tma_store( - tidx, - warp_idx, - acc_pipeline, - tiled_mma, - tma_atom_c, - tCtAcc_base, - sC, - tCgC, - epi_tile, - tile_sched, - epilogue_op, + c_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + 32 * len(self.epilogue_warp_id), ) + c_pipeline = pipeline.PipelineTmaStore.create( + num_stages=self.num_c_stage, producer_group=c_producer_group + ) + while work_tile.is_valid_tile: + # Get tile coord from tile scheduler + cur_tile_coord = work_tile.tile_idx + mma_tile_coord_mnl = ( + cur_tile_coord[0] // cute.size(tiled_mma.thr_id.shape), + cur_tile_coord[1], + cur_tile_coord[2], + ) + # + # Pre-advance to next tile + # + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + + num_tiles_executed = tile_sched.num_tiles_executed + if cutlass.const_expr(self.use_tma_store): + acc_consumer_state = utils.gemm.sm100.epilogue_tma_store( + self, + tidx, + warp_idx, + tma_atom_c, + tCtAcc_base, + sC, + tCgC, + epi_tile, + num_tiles_executed, + epilogue_op, + mma_tile_coord_mnl, + acc_consumer_state, + acc_pipeline, + c_pipeline, + ) + else: + acc_consumer_state = utils.gemm.sm100.epilogue( + self, + tidx, + tCtAcc_base, + tCgC, + epi_tile, + epilogue_op, + mma_tile_coord_mnl, + acc_consumer_state, + acc_pipeline, + ) + + if cutlass.const_expr(self.use_tma_store): + # Wait for C store complete + c_pipeline.producer_tail() else: - self.epilogue( - tidx, - acc_pipeline, - tiled_mma, - tCtAcc_base, - tCgC, - epi_tile, - tile_sched, - epilogue_op, - tmem_dealloc_barrier, - ) + # Synchronize before TMEM dealloc (done by the caller) + tmem_dealloc_barrier.arrive_and_wait() # # Dealloc the tensor memory buffer @@ -1011,354 +1045,6 @@ class PersistentDenseGemmKernel: tmem.relinquish_alloc_permit() tmem.free(tmem_ptr) - @cute.jit - def epilogue_tma_store( - self, - epi_tidx: cutlass.Int32, - warp_idx: cutlass.Int32, - acc_pipeline: pipeline.PipelineAsync, - tiled_mma: cute.TiledMma, - tma_atom_c: cute.CopyAtom, - # Input of epilogue - tCtAcc_base: cute.Tensor, - # Staging of epilogue - sC: cute.Tensor, - # Output of epilogue - tCgC: cute.Tensor, - epi_tile: cute.Tile, - tile_sched: utils.StaticPersistentTileScheduler, - epilogue_op: cutlass.Constexpr, - ) -> None: - tiled_copy_t2r, tTR_tAcc_base, tTR_rAcc = self.epilog_tmem_copy_and_partition( - epi_tidx, tCtAcc_base, tCgC, epi_tile, self.use_2cta_instrs - ) - - tTR_rC = cute.make_rmem_tensor(tTR_rAcc.shape, self.c_dtype) - tiled_copy_r2s, tRS_rC, tRS_sC = self.epilog_smem_copy_and_partition( - tiled_copy_t2r, tTR_rC, epi_tidx, sC - ) - - # (EPI_TILE_M, EPI_TILE_N, EPI_M, EPI_N, RestM, RestN, RestL) - tCgC_epi = cute.flat_divide( - tCgC[((None, None), 0, 0, None, None, None)], epi_tile - ) - # ((ATOM_V, REST_V), EPI_M, EPI_N) - # ((ATOM_V, REST_V), EPI_M, EPI_N, RestM, RestN, RestL) - bSG_sC, bSG_gC_partitioned = cpasync.tma_partition( - tma_atom_c, - 0, - cute.make_layout(1), - cute.group_modes(sC, 0, 2), - cute.group_modes(tCgC_epi, 0, 2), - ) - - acc_consumer_state = pipeline.make_pipeline_state( - pipeline.PipelineUserType.Consumer, self.num_acc_stage - ) - - # Threads/warps participating in tma store pipeline - c_producer_group = pipeline.CooperativeGroup( - pipeline.Agent.Thread, - 32 * len(self.epilog_warp_id), - ) - c_pipeline = pipeline.PipelineTmaStore.create( - num_stages=self.num_c_stage, producer_group=c_producer_group - ) - - epilog_sync_barrier = pipeline.NamedBarrier( - barrier_id=self.epilog_sync_bar_id, - num_threads=32 * len(self.epilog_warp_id), - ) - - work_tile = tile_sched.initial_work_tile_info() - while work_tile.is_valid_tile: - # Get tile coord from tile scheduler - cur_tile_coord = work_tile.tile_idx - mma_tile_coord_mnl = ( - cur_tile_coord[0] // cute.size(tiled_mma.thr_id.shape), - cur_tile_coord[1], - cur_tile_coord[2], - ) - - # - # Slice to per mma tile index - # - # ((ATOM_V, REST_V), EPI_M, EPI_N) - bSG_gC = bSG_gC_partitioned[(None, None, None, *mma_tile_coord_mnl)] - - # Set tensor memory buffer for current tile - # (T2R, T2R_M, T2R_N, EPI_M, EPI_M) - tTR_tAcc = tTR_tAcc_base[ - (None, None, None, None, None, acc_consumer_state.index) - ] - - # - # Wait for accumulator buffer full - # - acc_pipeline.consumer_wait(acc_consumer_state) - - tTR_tAcc = cute.group_modes(tTR_tAcc, 3, cute.rank(tTR_tAcc)) - bSG_gC = cute.group_modes(bSG_gC, 1, cute.rank(bSG_gC)) - - # - # Store accumulator to global memory in subtiles - # - subtile_cnt = cute.size(tTR_tAcc.shape, mode=[3]) - num_prev_subtiles = tile_sched.num_tiles_executed * subtile_cnt - for subtile_idx in cutlass.range(subtile_cnt): - # - # Load accumulator from tensor memory buffer to register - # - tTR_tAcc_mn = tTR_tAcc[(None, None, None, subtile_idx)] - cute.copy(tiled_copy_t2r, tTR_tAcc_mn, tTR_rAcc) - - # - # Convert to C type - # - acc_vec = tiled_copy_r2s.retile(tTR_rAcc).load() - acc_vec = epilogue_op(acc_vec.to(self.c_dtype)) - tRS_rC.store(acc_vec) - - # - # Store C to shared memory - # - c_buffer = (num_prev_subtiles + subtile_idx) % self.num_c_stage - cute.copy(tiled_copy_r2s, tRS_rC, tRS_sC[(None, None, None, c_buffer)]) - # Fence and barrier to make sure shared memory store is visible to TMA store - cute.arch.fence_proxy("async.shared", space="cta") - epilog_sync_barrier.arrive_and_wait() - - # - # TMA store C to global memory - # - if warp_idx == self.epilog_warp_id[0]: - cute.copy( - tma_atom_c, - bSG_sC[(None, c_buffer)], - bSG_gC[(None, subtile_idx)], - ) - # Fence and barrier to make sure shared memory store is visible to TMA store - c_pipeline.producer_commit() - c_pipeline.producer_acquire() - epilog_sync_barrier.arrive_and_wait() - - epilog_sync_barrier.arrive_and_wait() - - # - # Async arrive accumulator buffer empty - # - with cute.arch.elect_one(): - acc_pipeline.consumer_release(acc_consumer_state) - acc_consumer_state.advance() - - # - # Advance to next tile - # - tile_sched.advance_to_next_work() - work_tile = tile_sched.get_current_work() - - # Wait for C store complete - c_pipeline.producer_tail() - - @cute.jit - def epilogue( - self, - epi_tidx: cutlass.Int32, - acc_pipeline: pipeline.PipelineAsync, - tiled_mma: cute.TiledMma, - tCtAcc_base: cute.Tensor, - tCgC: cute.Tensor, - epi_tile: cute.Tile, - tile_sched: utils.StaticPersistentTileScheduler, - epilogue_op: cutlass.Constexpr, - tmem_dealloc_barrier: pipeline.NamedBarrier, - ) -> None: - tiled_copy_t2r, tTR_tAcc_base, tTR_rAcc = self.epilog_tmem_copy_and_partition( - epi_tidx, tCtAcc_base, tCgC, epi_tile, self.use_2cta_instrs - ) - - gC_epi = cute.flat_divide( - tCgC[((None, None), 0, 0, None, None, None)], epi_tile - ) - # (T2R, T2R_M, T2R_N, EPI_M, EPI_N, RestM, RestN, RestL) - thr_copy_t2r = tiled_copy_t2r.get_slice(epi_tidx) - tTR_gC_partitioned = thr_copy_t2r.partition_D(gC_epi) - # (T2R, T2R_M, T2R_N) - tTR_rC = cute.make_rmem_tensor( - tTR_gC_partitioned[(None, None, None, 0, 0, 0, 0, 0)].shape, self.c_dtype - ) - simt_atom = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), self.c_dtype) - - acc_consumer_state = pipeline.make_pipeline_state( - pipeline.PipelineUserType.Consumer, self.num_acc_stage - ) - - work_tile = tile_sched.initial_work_tile_info() - while work_tile.is_valid_tile: - # Get tile coord from tile scheduler - cur_tile_coord = work_tile.tile_idx - mma_tile_coord_mnl = ( - cur_tile_coord[0] // cute.size(tiled_mma.thr_id.shape), - cur_tile_coord[1], - cur_tile_coord[2], - ) - - # - # Slice to per mma tile index - # - # (T2R, T2R_M, T2R_N, EPI_M, EPI_N) - tTR_gC = tTR_gC_partitioned[ - (None, None, None, None, None, *mma_tile_coord_mnl) - ] - - # Set tensor memory buffer for current tile - # (T2R, T2R_M, T2R_N, EPI_M, EPI_N) - tTR_tAcc = tTR_tAcc_base[ - (None, None, None, None, None, acc_consumer_state.index) - ] - - tTR_tAcc = cute.group_modes(tTR_tAcc, 3, cute.rank(tTR_tAcc)) - tTR_gC = cute.group_modes(tTR_gC, 3, cute.rank(tTR_gC)) - - # - # Wait for accumulator buffer full - # - acc_pipeline.consumer_wait(acc_consumer_state) - - # - # Store accumulator to global memory in subtiles - # - subtile_cnt = cute.size(tTR_tAcc.shape, mode=[3]) - for subtile_idx in cutlass.range(subtile_cnt): - # - # Load accumulator from tensor memory buffer to register - # - tTR_tAcc_mn = tTR_tAcc[(None, None, None, subtile_idx)] - cute.copy(tiled_copy_t2r, tTR_tAcc_mn, tTR_rAcc) - - # - # Convert to C type - # - acc_vec = tTR_rAcc.load() - acc_vec = epilogue_op(acc_vec.to(self.c_dtype)) - tTR_rC.store(acc_vec) - - # - # Store C to global memory - # - cute.copy(simt_atom, tTR_rC, tTR_gC[(None, None, None, subtile_idx)]) - - # - # Async arrive accumulator buffer empty - # - with cute.arch.elect_one(): - acc_pipeline.consumer_release(acc_consumer_state) - acc_consumer_state.advance() - - # Advance to next tile - tile_sched.advance_to_next_work() - work_tile = tile_sched.get_current_work() - - # Synchronize before TMEM dealloc (done by the caller) - tmem_dealloc_barrier.arrive_and_wait() - - def epilog_tmem_copy_and_partition( - self, - tidx: cutlass.Int32, - tAcc: cute.Tensor, - gC_mnl: cute.Tensor, - epi_tile: cute.Tile, - use_2cta_instrs: Union[cutlass.Boolean, bool], - ) -> Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]: - """ - Make tiledCopy for tensor memory load, then use it to partition tensor memory (source) and register array (destination). - - :param tidx: The thread index in epilogue warp groups - :type tidx: cutlass.Int32 - :param tAcc: The accumulator tensor to be copied and partitioned - :type tAcc: cute.Tensor - :param gC_mnl: The global tensor C - :type gC_mnl: cute.Tensor - :param epi_tile: The epilogue tiler - :type epi_tile: cute.Tile - :param use_2cta_instrs: Whether use_2cta_instrs is enabled - :type use_2cta_instrs: bool - - :return: A tuple containing (tiled_copy_t2r, tTR_tAcc, tTR_rAcc) where: - - tiled_copy_t2r: The tiled copy operation for tmem to register copy(t2r) - - tTR_tAcc: The partitioned accumulator tensor - - tTR_rAcc: The accumulated tensor in register used to hold t2r results - :rtype: Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor] - """ - # Make tiledCopy for tensor memory load - copy_atom_t2r = utils.sm100.get_tmem_load_op( - self.cta_tile_shape_mnk, - self.c_layout, - self.c_dtype, - self.acc_dtype, - epi_tile, - use_2cta_instrs, - ) - # (EPI_TILE_M, EPI_TILE_N, EPI_M, EPI_N, STAGE) - tAcc_epi = cute.flat_divide(tAcc[((None, None), 0, 0, None)], epi_tile) - # (EPI_TILE_M, EPI_TILE_N) - tiled_copy_t2r = tcgen05.make_tmem_copy( - copy_atom_t2r, tAcc_epi[(None, None, 0, 0, 0)] - ) - - thr_copy_t2r = tiled_copy_t2r.get_slice(tidx) - # (T2R, T2R_M, T2R_N, EPI_M, EPI_M, STAGE) - tTR_tAcc = thr_copy_t2r.partition_S(tAcc_epi) - - # (EPI_TILE_M, EPI_TILE_N, EPI_M, EPI_N, RestM, RestN, RestL) - gC_mnl_epi = cute.flat_divide( - gC_mnl[((None, None), 0, 0, None, None, None)], epi_tile - ) - # (T2R, T2R_M, T2R_N, EPI_M, EPI_N, RestM, RestN, RestL) - tTR_gC = thr_copy_t2r.partition_D(gC_mnl_epi) - # (T2R, T2R_M, T2R_N) - tTR_rAcc = cute.make_rmem_tensor( - tTR_gC[(None, None, None, 0, 0, 0, 0, 0)].shape, self.acc_dtype - ) - return tiled_copy_t2r, tTR_tAcc, tTR_rAcc - - def epilog_smem_copy_and_partition( - self, - tiled_copy_t2r: cute.TiledCopy, - tTR_rC: cute.Tensor, - tidx: cutlass.Int32, - sC: cute.Tensor, - ) -> Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]: - """ - Make tiledCopy for shared memory store, then use it to partition register array (source) and shared memory (destination). - - :param tiled_copy_t2r: The tiled copy operation for tmem to register copy(t2r) - :type tiled_copy_t2r: cute.TiledCopy - :param tTR_rC: The partitioned accumulator tensor - :type tTR_rC: cute.Tensor - :param tidx: The thread index in epilogue warp groups - :type tidx: cutlass.Int32 - :param sC: The shared memory tensor to be copied and partitioned - :type sC: cute.Tensor - :type sepi: cute.Tensor - - :return: A tuple containing (tiled_copy_r2s, tRS_rC, tRS_sC) where: - - tiled_copy_r2s: The tiled copy operation for register to smem copy(r2s) - - tRS_rC: The partitioned tensor C (register source) - - tRS_sC: The partitioned tensor C (smem destination) - :rtype: Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor] - """ - copy_atom_r2s = utils.sm100.get_smem_store_op( - self.c_layout, self.c_dtype, self.acc_dtype, tiled_copy_t2r - ) - tiled_copy_r2s = cute.make_tiled_copy_D(copy_atom_r2s, tiled_copy_t2r) - # (R2S, R2S_M, R2S_N, PIPE_D) - thr_copy_r2s = tiled_copy_r2s.get_slice(tidx) - tRS_sC = thr_copy_r2s.partition_D(sC) - # (R2S, R2S_M, R2S_N) - tRS_rC = tiled_copy_r2s.retile(tTR_rC) - return tiled_copy_r2s, tRS_rC, tRS_sC - @staticmethod def _compute_grid( c: cute.Tensor, @@ -1401,6 +1087,7 @@ class PersistentDenseGemmKernel: tiled_mma: cute.TiledMma, mma_tiler: Tuple[int, int, int], num_acc_stage: int, + arch: str, ) -> int: """ Compute the number of tensor memory allocation columns. @@ -1417,25 +1104,29 @@ class PersistentDenseGemmKernel: """ acc_shape = tiled_mma.partition_shape_C(mma_tiler[:2]) tCtAcc_fake = tiled_mma.make_fragment_C(cute.append(acc_shape, num_acc_stage)) - num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols(tCtAcc_fake) + num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols(tCtAcc_fake, arch=arch) return num_tmem_alloc_cols - def is_valid_dtypes( - self, ab_dtype: Type[cutlass.Numeric], c_dtype: Type[cutlass.Numeric] - ) -> bool: + def check_supported_dtypes( + self, + a_dtype: Type[cutlass.Numeric], + b_dtype: Type[cutlass.Numeric], + c_dtype: Type[cutlass.Numeric], + ): """ Check if the dtypes are valid - :param ab_dtype: The data type of the A and B operands - :type ab_dtype: Type[cutlass.Numeric] + :param a_dtype: The data type of the A operands + :type a_dtype: Type[cutlass.Numeric] + :param b_dtype: The data type of the B operands + :type b_dtype: Type[cutlass.Numeric] :param acc_dtype: The data type of the accumulator :type acc_dtype: Type[cutlass.Numeric] :param c_dtype: The data type of the output tensor :type c_dtype: Type[cutlass.Numeric] - :return: True if the dtypes are valid, False otherwise - :rtype: bool + :raises testing.CantImplementError: If the dtypes are invalid """ valid_ab_dtypes = { cutlass.Float16, @@ -1446,11 +1137,15 @@ class PersistentDenseGemmKernel: cutlass.Float8E4M3FN, cutlass.Float8E5M2, } - if ab_dtype not in valid_ab_dtypes: - return False + if a_dtype not in valid_ab_dtypes or b_dtype not in valid_ab_dtypes: + raise testing.CantImplementError( + f"Unsupported AB dtype: {a_dtype} and {b_dtype}" + ) if self.acc_dtype not in {cutlass.Float32, cutlass.Float16, cutlass.Int32}: - return False + raise testing.CantImplementError( + f"Unsupported accumulator dtype: {self.acc_dtype}" + ) # Define compatibility mapping between accumulator type and AB type acc_ab_compatibility = { @@ -1469,8 +1164,13 @@ class PersistentDenseGemmKernel: cutlass.Int32: {cutlass.Uint8, cutlass.Int8}, } # Check compatibility between accumulator type and AB type - if ab_dtype not in acc_ab_compatibility[self.acc_dtype]: - return False + if ( + a_dtype not in acc_ab_compatibility[self.acc_dtype] + or b_dtype not in acc_ab_compatibility[self.acc_dtype] + ): + raise testing.CantImplementError( + f"Unsupported AB dtype: {a_dtype} and {b_dtype} for accumulator dtype: {self.acc_dtype}" + ) # Define compatibility mapping between accumulator type and C type acc_c_compatibility = { @@ -1499,28 +1199,32 @@ class PersistentDenseGemmKernel: } # Check compatibility between accumulator type and C type if c_dtype not in acc_c_compatibility[self.acc_dtype]: - return False + raise testing.CantImplementError( + f"Unsupported C dtype: {c_dtype} for accumulator dtype: {self.acc_dtype}" + ) - return True - - def is_valid_mma_tiler_and_cluster_shape(self) -> bool: + def check_mma_tiler_and_cluster_shape(self): """Check if the mma tiler and cluster shape are valid. - :return: True if the mma tiler and cluster shape are valid, False otherwise - :rtype: bool + :raises testing.CantImplementError: If the mma tiler and cluster shape are invalid """ - is_valid = True # Skip invalid mma tile shape if not ( (not self.use_2cta_instrs and self.mma_tiler_mn[0] in [64, 128]) or (self.use_2cta_instrs and self.mma_tiler_mn[0] in [128, 256]) ): - is_valid = False + raise testing.CantImplementError( + f"Invalid mma tiler & use_2cta_instrs: {self.mma_tiler_mn}, {self.use_2cta_instrs}" + ) if self.mma_tiler_mn[1] not in range(32, 257, 32): - is_valid = False + raise testing.CantImplementError( + f"Invalid mma tiler N: {self.mma_tiler_mn[1]}" + ) # Skip illegal cluster shape if self.cluster_shape_mn[0] % (2 if self.use_2cta_instrs else 1) != 0: - is_valid = False + raise testing.CantImplementError( + f"Invalid cluster shape M: {self.cluster_shape_mn[0]}" + ) # Skip invalid cluster shape is_power_of_2 = lambda x: x > 0 and (x & (x - 1)) == 0 if ( @@ -1530,21 +1234,23 @@ class PersistentDenseGemmKernel: or not is_power_of_2(self.cluster_shape_mn[0]) or not is_power_of_2(self.cluster_shape_mn[1]) ): - is_valid = False - return is_valid + raise testing.CantImplementError( + f"Invalid cluster shape: {self.cluster_shape_mn}" + ) - def is_valid_tensor_alignment( + def check_tensor_alignment( self, m: int, n: int, k: int, l: int, - ab_dtype: Type[cutlass.Numeric], + a_dtype: Type[cutlass.Numeric], + b_dtype: Type[cutlass.Numeric], c_dtype: Type[cutlass.Numeric], a_major: str, b_major: str, c_major: str, - ) -> bool: + ): """ Check if the tensor alignment is valid @@ -1556,8 +1262,10 @@ class PersistentDenseGemmKernel: :type k: int :param l: The number of columns in the C tensor :type l: int - :param ab_dtype: The data type of the A and B operands - :type ab_dtype: Type[cutlass.Numeric] + :param a_dtype: The data type of the A operands + :type a_dtype: Type[cutlass.Numeric] + :param b_dtype: The data type of the B operands + :type b_dtype: Type[cutlass.Numeric] :param c_dtype: The data type of the output tensor :type c_dtype: Type[cutlass.Numeric] :param a_major: The major axis of the A tensor @@ -1567,10 +1275,8 @@ class PersistentDenseGemmKernel: :param c_major: The major axis of the C tensor :type c_major: str - :return: True if the problem shape is valid, False otherwise - :rtype: bool + :raises testing.CantImplementError: If the tensor alignment is invalid """ - is_valid = True # TODO: move to utils def check_contiguous_16B_alignment(dtype, is_mode0_major, tensor_shape): @@ -1580,14 +1286,15 @@ class PersistentDenseGemmKernel: return num_major_elements % num_contiguous_elements == 0 if ( - not check_contiguous_16B_alignment(ab_dtype, a_major == "m", (m, k, l)) - or not check_contiguous_16B_alignment(ab_dtype, b_major == "n", (n, k, l)) + not check_contiguous_16B_alignment(a_dtype, a_major == "m", (m, k, l)) + or not check_contiguous_16B_alignment(b_dtype, b_major == "n", (n, k, l)) or not check_contiguous_16B_alignment(c_dtype, c_major == "m", (m, n, l)) ): - is_valid = False - return is_valid + raise testing.CantImplementError( + f"Invalid tensor alignment: {m}, {n}, {k}, {l}, {a_dtype}, {b_dtype}, {c_dtype}, {a_major}, {b_major}, {c_major}" + ) - def is_valid_epilog_store_option(self, m: int, n: int) -> bool: + def check_epilog_store_option(self, m: int, n: int): """ Check if the epilogue store option is valid @@ -1596,11 +1303,8 @@ class PersistentDenseGemmKernel: :param n: The number of columns in the B tensor :type n: int - :return: True if the epilogue store option is valid, False otherwise - :rtype: bool + :raises testing.CantImplementError: If the epilogue store option is invalid """ - - is_valid = True # None TMA store version does not have predication, can not support OOB tiles cta_tile_shape_mn = ( self.mma_tiler_mn[0] // (2 if self.use_2cta_instrs else 1), @@ -1608,13 +1312,15 @@ class PersistentDenseGemmKernel: ) if not self.use_tma_store: if not (m % cta_tile_shape_mn[0] == 0 and n % cta_tile_shape_mn[1] == 0): - is_valid = False - return is_valid + raise testing.CantImplementError( + f"Invalid epilog store option: {m}, {n}" + ) def can_implement( self, mnkl: Tuple[int, int, int, int], - ab_dtype: Type[cutlass.Numeric], + a_dtype: Type[cutlass.Numeric], + b_dtype: Type[cutlass.Numeric], c_dtype: Type[cutlass.Numeric], a_major: str, b_major: str, @@ -1625,8 +1331,10 @@ class PersistentDenseGemmKernel: :param mnkl: Problem size as a tuple (M, N, K, L). :type mnkl: Tuple[int, int, int, int] - :param ab_dtype: Data type for input tensors A and B. - :type ab_dtype: Type[cutlass.Numeric] + :param a_dtype: Data type for input tensors A. + :type a_dtype: Type[cutlass.Numeric] + :param b_dtype: Data type for input tensors B. + :type b_dtype: Type[cutlass.Numeric] :param c_dtype: Data type for output tensor C. :type c_dtype: Type[cutlass.Numeric] :param a_major: Major dimension of the A tensor layout ("m" or "k"). @@ -1639,25 +1347,20 @@ class PersistentDenseGemmKernel: :rtype: bool """ - # Skip unsupported types - if not self.is_valid_dtypes(ab_dtype, c_dtype): - return False + try: + # Skip unsupported types + self.check_supported_dtypes(a_dtype, b_dtype, c_dtype) - # Skip invalid mma tile shape and cluster shape - if not self.is_valid_mma_tiler_and_cluster_shape(): - return False + # Skip invalid mma tile shape and cluster shape + self.check_mma_tiler_and_cluster_shape() - # Unpack mnkl for clarity in calling the epilog check - m, n, k, l = mnkl - # Skip illegal problem shape for load/store alignment - if not self.is_valid_tensor_alignment( - m, n, k, l, ab_dtype, c_dtype, a_major, b_major, c_major - ): + m, n, k, l = mnkl + self.check_tensor_alignment( + m, n, k, l, a_dtype, b_dtype, c_dtype, a_major, b_major, c_major + ) + self.check_epilog_store_option(m, n) + except testing.CantImplementError: return False - # Skip invalid epilogue store option - if not self.is_valid_epilog_store_option(m, n): - return False - return True @@ -1702,45 +1405,117 @@ def bmm( gemm_op(a, b, c, max_active_clusters, stream, epilogue_op) +@lru_cache(maxsize=1) def prepare_tensors( mnkl: Tuple[int, int, int, int], - ab_dtype: Type[cutlass.Numeric], + a_dtype: Type[cutlass.Numeric], + b_dtype: Type[cutlass.Numeric], c_dtype: Type[cutlass.Numeric], a_major: str, b_major: str, c_major: str, init_random: bool = True, + normal_mean: float = 0.0, + normal_std: float = 1.0, ): + """Prepare tensors for GEMM. + + Returns: + Tuple of (a_f32, b_f32, c_f32, a_storage, b_storage, c_storage): + - *_f32: Float32 tensors with the logical data (for reference and fp8 conversion) + - *_storage: Storage tensors for DLPack (uint8 for fp8, otherwise the target dtype) + """ import torch from cutlass.torch import dtype as torch_dtype m, n, k, l = mnkl if a_major == "k": - a = torch.empty((l, m, k), dtype=torch.float32, device="cuda") + a_f32 = torch.empty((l, m, k), dtype=torch.float32, device="cuda") elif a_major == "m": - a = torch.empty((l, k, m), dtype=torch.float32, device="cuda").permute(0, 2, 1) + a_f32 = torch.empty((l, k, m), dtype=torch.float32, device="cuda").permute( + 0, 2, 1 + ) if b_major == "n": - b = torch.empty((l, k, n), dtype=torch.float32, device="cuda") + b_f32 = torch.empty((l, k, n), dtype=torch.float32, device="cuda") elif b_major == "k": - b = torch.empty((l, n, k), dtype=torch.float32, device="cuda").permute(0, 2, 1) + b_f32 = torch.empty((l, n, k), dtype=torch.float32, device="cuda").permute( + 0, 2, 1 + ) if c_major == "n": - c = torch.empty((l, m, n), dtype=torch.float32, device="cuda") + c_f32 = torch.empty((l, m, n), dtype=torch.float32, device="cuda") elif c_major == "m": - c = torch.empty((l, n, m), dtype=torch.float32, device="cuda").permute(0, 2, 1) + c_f32 = torch.empty((l, n, m), dtype=torch.float32, device="cuda").permute( + 0, 2, 1 + ) if init_random: - a.random_(-2, 3) - b.random_(-2, 3) - c.random_(-2, 3) + # Uniform random initialization in range [-2, 3) + a_f32.random_(-2, 3) + b_f32.random_(-2, 3) + c_f32.random_(-2, 3) - return ( - a.to(dtype=torch_dtype(ab_dtype)), - b.to(dtype=torch_dtype(ab_dtype)), - c.to(dtype=torch_dtype(c_dtype)), + else: + # Normal (Gaussian) initialization with user-specified mean and std + a_f32.normal_(mean=normal_mean, std=normal_std) + b_f32.normal_(mean=normal_mean, std=normal_std) + c_f32.normal_(mean=normal_mean, std=normal_std) + + # For float8 types, use uint8 as storage type to avoid dlpack limitation + # (dlpack doesn't support float8 types) + # For other types, convert to the target dtype + a_storage_dtype = torch.uint8 if is_fp8_dtype(a_dtype) else torch_dtype(a_dtype) + b_storage_dtype = torch.uint8 if is_fp8_dtype(b_dtype) else torch_dtype(b_dtype) + c_storage_dtype = torch.uint8 if is_fp8_dtype(c_dtype) else torch_dtype(c_dtype) + + a_storage = a_f32.to(dtype=a_storage_dtype) + b_storage = b_f32.to(dtype=b_storage_dtype) + c_storage = c_f32.to(dtype=c_storage_dtype) + + return (a_f32, b_f32, c_f32, a_storage, b_storage, c_storage) + + +@lru_cache(maxsize=1) +def compile_bmm( + mnkl: Tuple[int, int, int, int], + a: cute.Tensor, + b: cute.Tensor, + c: cute.Tensor, + acc_dtype: Type[cutlass.Numeric], + a_major: str, + b_major: str, + c_major: str, + mma_tiler_mn: Tuple[int, int] = (256, 256), + cluster_shape_mn: Tuple[int, int] = (2, 1), + max_active_clusters: cutlass.Constexpr = None, + use_2cta_instrs: bool = True, + use_tma_store: bool = True, + epilogue_op: cutlass.Constexpr = lambda x: x, +): + from cutlass.cute.runtime import make_fake_stream + + gemm = PersistentDenseGemmKernel( + acc_dtype, + use_2cta_instrs, + mma_tiler_mn, + cluster_shape_mn, + use_tma_store, ) + # Check if configuration can be implemented + can_implement = gemm.can_implement( + mnkl, a.element_type, b.element_type, c.element_type, a_major, b_major, c_major + ) + if not can_implement: + raise testing.CantImplementError( + f"The current config which is invalid/unsupported: use_2cta_instrs = {use_2cta_instrs}, " + f"mma_tiler_mn = {mma_tiler_mn}, cluster_shape_mn = {cluster_shape_mn}, " + f"use_tma_store = {use_tma_store}" + ) + + stream = make_fake_stream() + return cute.compile(bmm, gemm, a, b, c, max_active_clusters, stream, epilogue_op) def run( @@ -1808,36 +1583,9 @@ def run( :return: Execution time of the GEMM kernel. :rtype: float """ - print("Running Blackwell Persistent Dense GEMM test with:") - print(f"mnkl: {mnkl}") - print(f"AB dtype: {ab_dtype}, C dtype: {c_dtype}, Acc dtype: {acc_dtype}") - print(f"Matrix majors - A: {a_major}, B: {b_major}, C: {c_major}") - print(f"Mma Tiler (M, N): {mma_tiler_mn}, Cluster Shape (M, N): {cluster_shape_mn}") - print(f"2CTA MMA instructions: {'True' if use_2cta_instrs else 'False'}") - print(f"Use TMA Store: {'True' if use_tma_store else 'False'}") - print(f"Tolerance: {tolerance}") - print(f"Warmup iterations: {warmup_iterations}") - print(f"Iterations: {iterations}") - print(f"Skip reference checking: {skip_ref_check}") - print(f"Use cold L2: {'True' if use_cold_l2 else 'False'}") - import torch from cutlass.torch import dtype as torch_dtype - # Build GEMM object - gemm_op = PersistentDenseGemmKernel( - acc_dtype, use_2cta_instrs, mma_tiler_mn, cluster_shape_mn, use_tma_store - ) - can_implement = gemm_op.can_implement( - mnkl, ab_dtype, c_dtype, a_major, b_major, c_major - ) - if not can_implement: - raise testing.CantImplementError( - f"The current config which is invalid/unsupported: use_2cta_instrs = {use_2cta_instrs}, " - f"mma_tiler_mn = {mma_tiler_mn}, cluster_shape_mn = {cluster_shape_mn}, " - f"use_tma_store = {use_tma_store}" - ) - if not torch.cuda.is_available(): raise RuntimeError("GPU is required to run this example!") @@ -1852,67 +1600,96 @@ def run( ) # Run and verify BMM with torch - a, b, c = prepare_tensors(mnkl, ab_dtype, c_dtype, a_major, b_major, c_major) + a_f32, b_f32, c_f32, a_storage, b_storage, c_storage = prepare_tensors( + mnkl, ab_dtype, ab_dtype, c_dtype, a_major, b_major, c_major + ) - # Leading dim is 2 leading_dim_a = 2 if a_major == "k" else 1 leading_dim_b = 1 if b_major == "k" else 2 leading_dim_c = 2 if c_major == "n" else 1 - a_ = from_dlpack(a).mark_layout_dynamic(leading_dim=leading_dim_a) - b_ = from_dlpack(b).mark_layout_dynamic(leading_dim=leading_dim_b) - c_ = from_dlpack(c).mark_layout_dynamic(leading_dim=leading_dim_c) + # Create CuTe tensors, passing float32 source for fp8 conversion + a_ = create_cute_tensor_for_fp8( + a_storage, ab_dtype, leading_dim_a, source_f32_tensor=a_f32 + ) + b_ = create_cute_tensor_for_fp8( + b_storage, ab_dtype, leading_dim_b, source_f32_tensor=b_f32 + ) + c_ = create_cute_tensor_for_fp8( + c_storage, c_dtype, leading_dim_c, source_f32_tensor=c_f32 + ) - compiled_fn = cute.compile( - bmm, - gemm_op, + compiled_fn = compile_bmm( + mnkl, a_, b_, c_, + acc_dtype, + a_major, + b_major, + c_major, + mma_tiler_mn, + cluster_shape_mn, max_active_clusters, - current_stream, + use_2cta_instrs, + use_tma_store, epilogue_op=lambda x: x, ) + print("Running Blackwell Persistent Dense GEMM test with:") + print(f"mnkl: {mnkl}") + print(f"Tolerance: {tolerance}") + print(f"Warmup iterations: {warmup_iterations}") + print(f"Iterations: {iterations}") + print(f"Skip reference checking: {skip_ref_check}") + print(f"Use cold L2: {'True' if use_cold_l2 else 'False'}") + if not skip_ref_check: # Use small random number for deterministic result for reference check compiled_fn(a_, b_, c_, current_stream) # Manually quantize to be comparable + # Use float32 source data for reference calculation ref = ( - torch.bmm(a.to(dtype=torch.float32), b.to(dtype=torch.float32)) + torch.bmm(a_f32, b_f32) .to(dtype=torch_dtype(c_dtype)) .to(dtype=torch.float32) ) + # Read back the result from CuTe tensor (c_storage was updated in-place) torch.testing.assert_close( - c.to(dtype=torch.float32), ref, atol=tolerance, rtol=1e-03 + c_storage.to(dtype=torch.float32), ref, atol=tolerance, rtol=1e-03 ) if not benchmark: return 0 def generate_tensors(): - init_normal = ab_dtype not in [cutlass.Int8, cutlass.Uint8] - a, b, c = prepare_tensors( + a_f32, b_f32, c_f32, a_st, b_st, c_st = prepare_tensors( mnkl, ab_dtype, + ab_dtype, c_dtype, a_major, b_major, c_major, - init_random=not init_normal, ) - a_ = from_dlpack(a).mark_layout_dynamic(leading_dim=leading_dim_a) - b_ = from_dlpack(b).mark_layout_dynamic(leading_dim=leading_dim_b) - c_ = from_dlpack(c).mark_layout_dynamic(leading_dim=leading_dim_c) + a_ = create_cute_tensor_for_fp8( + a_st, ab_dtype, leading_dim_a, source_f32_tensor=a_f32 + ) + b_ = create_cute_tensor_for_fp8( + b_st, ab_dtype, leading_dim_b, source_f32_tensor=b_f32 + ) + c_ = create_cute_tensor_for_fp8( + c_st, c_dtype, leading_dim_c, source_f32_tensor=c_f32 + ) return testing.JitArguments(a_, b_, c_, current_stream) workspace_count = 1 if use_cold_l2: one_workspace_bytes = ( - a.numel() * a.element_size() - + b.numel() * b.element_size() - + c.numel() * c.element_size() + a_storage.numel() * a_storage.element_size() + + b_storage.numel() * b_storage.element_size() + + c_storage.numel() * c_storage.element_size() ) workspace_count = testing.get_workspace_count( one_workspace_bytes, warmup_iterations, iterations @@ -1929,6 +1706,11 @@ def run( ) +def compute_tflops(time_ns, m, n, k): + return 2.0 * m * n * k / time_ns / 1000.0 + + + def _parse_comma_separated_ints(s: str) -> Tuple[int, ...]: try: return tuple(int(x.strip()) for x in s.split(",")) @@ -1939,7 +1721,6 @@ def _parse_comma_separated_ints(s: str) -> Tuple[int, ...]: def prepare_parser(): - parser = argparse.ArgumentParser( description="Example of Dense Persistent GEMM on Blackwell." ) @@ -1974,7 +1755,14 @@ def prepare_parser(): "--tolerance", type=float, default=1e-01, help="Tolerance for validation" ) parser.add_argument( - "--benchmark", action="store_true", help="Only benchmark the kernel" + "--benchmark", + type=str, + default="default", + choices=[ + "default", + "none", + ], + help="Benchmark the kernel with nsight or default (cute.testing.benchmark) or none", ) parser.add_argument( "--warmup_iterations", type=int, default=0, help="Warmup iterations" @@ -2018,6 +1806,20 @@ if __name__ == "__main__": if len(args.cluster_shape_mn) != 2: parser.error("--cluster_shape_mn must contain exactly 2 values") + print(f"[DSL INFO] Compiling Blackwell Persistent Dense GEMM with:") + print( + f"[DSL INFO] A dtype: {args.ab_dtype}, B dtype: {args.c_dtype}, C dtype: {args.acc_dtype}, Acc dtype: {args.acc_dtype}" + ) + print( + f"[DSL INFO] Matrix majors - A: {args.a_major}, B: {args.b_major}, C: {args.c_major}" + ) + print(f"[DSL INFO] Mma Tiler (M, N): {args.mma_tiler_mn}") + print(f"[DSL INFO] Cluster Shape (M, N): {args.cluster_shape_mn}") + print( + f"[DSL INFO] 2CTA MMA instructions: {'True' if args.use_2cta_instrs else 'False'}" + ) + print(f"[DSL INFO] Use TMA Store: {'True' if args.use_tma_store else 'False'}") + run( args.mnkl, args.ab_dtype, @@ -2035,6 +1837,6 @@ if __name__ == "__main__": args.iterations, args.skip_ref_check, args.use_cold_l2, - args.benchmark, + args.benchmark == "default", ) print("PASS") diff --git a/examples/python/CuTeDSL/blackwell/dense_gemm_persistent_dynamic.py b/examples/python/CuTeDSL/blackwell/dense_gemm_persistent_dynamic.py index 03d4b73f..7a259295 100644 --- a/examples/python/CuTeDSL/blackwell/dense_gemm_persistent_dynamic.py +++ b/examples/python/CuTeDSL/blackwell/dense_gemm_persistent_dynamic.py @@ -401,6 +401,7 @@ class PersistentDenseGemmKernel: # Setup clc stage by default self.num_clc_stage = 1 + assert self.num_clc_stage == 1, "Only single-stage CLC pipeline is supported" # Compute A/B/C shared memory layout self.a_smem_layout_staged = utils.sm100.make_smem_layout_a( @@ -615,8 +616,8 @@ class PersistentDenseGemmKernel: ] tmem_dealloc_mbar_ptr: cutlass.Int64 tmem_holding_buf: cutlass.Int32 - clc_ptr: cute.struct.MemRange[cutlass.Int64, self.num_clc_stage * 2] - clc_response_ptr: cute.struct.MemRange[cutlass.Int32, 1] + clc_mbar_ptr: cute.struct.MemRange[cutlass.Int64, 2] + clc_response: cute.struct.MemRange[cutlass.Int32, 4] smem = utils.SmemAllocator() storage = smem.allocate(SharedStorage) @@ -664,7 +665,7 @@ class PersistentDenseGemmKernel: pipeline.Agent.Thread, num_clc_consumer_threads ) clc_pipeline = pipeline.PipelineClcFetchAsync.create( - barrier_storage=storage.clc_ptr.data_ptr(), + barrier_storage=storage.clc_mbar_ptr.data_ptr(), num_stages=self.num_clc_stage, producer_group=clc_pipeline_producer_group, consumer_group=clc_pipeline_consumer_group, @@ -696,7 +697,7 @@ class PersistentDenseGemmKernel: pipeline_init_arrive(cluster_shape_mn=cluster_layout_vmnk, is_relaxed=True) # Initial clc response pointer - clc_response_ptr = storage.clc_response_ptr.data_ptr() + clc_response_ptr = storage.clc_response.data_ptr() clc_consumer_state = pipeline.make_pipeline_state( pipeline.PipelineUserType.Consumer, self.num_clc_stage @@ -1733,7 +1734,6 @@ def run( return 0 def generate_tensors(): - init_normal = ab_dtype not in [cutlass.Int8, cutlass.Uint8] a_f32, b_f32, c_f32, a_st, b_st, c_st = prepare_tensors( mnkl, ab_dtype, @@ -1742,7 +1742,6 @@ def run( a_major, b_major, c_major, - init_random=not init_normal, ) a_ = create_cute_tensor_for_fp8( a_st, ab_dtype, leading_dim_a, source_f32_tensor=a_f32 diff --git a/examples/python/CuTeDSL/blackwell/dense_gemm_software_pipeline.py b/examples/python/CuTeDSL/blackwell/dense_gemm_software_pipeline.py index 3952cd46..3622b211 100644 --- a/examples/python/CuTeDSL/blackwell/dense_gemm_software_pipeline.py +++ b/examples/python/CuTeDSL/blackwell/dense_gemm_software_pipeline.py @@ -30,15 +30,12 @@ import argparse from typing import Optional, Type, Tuple, Union import cuda.bindings.driver as cuda -import torch - import cutlass import cutlass.cute as cute import cutlass.utils as utils import cutlass.pipeline as pipeline from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait from cutlass.cute.nvgpu import cpasync, tcgen05 -import cutlass.torch as cutlass_torch import cutlass.utils.blackwell_helpers as sm100_utils import cutlass.cute.testing as testing @@ -767,6 +764,14 @@ class DenseGemmKernel: acc_pipeline.consumer_wait(acc_consumer_state) if cutlass.const_expr(self.use_tma_store): + # TODO: segment fault if I put here + # sC = smem.allocate_tensor( + # element_type=self.c_dtype, + # layout=c_smem_layout_staged.outer, + # byte_alignment=128, + # swizzle=c_smem_layout_staged.inner, + # ) + assert tma_atom_c is not None and sC is not None self.epilogue_tma_store( tidx, @@ -983,7 +988,10 @@ class DenseGemmKernel: c_buffer = subtile_idx % self.num_c_stage cute.copy(tiled_copy_r2s, tRS_rC, tRS_sC[(None, None, None, c_buffer)]) # Fence and barrier to make sure shared memory store is visible to TMA store - cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) pipeline.sync(barrier_id=1) # TMA store C to global memory @@ -1214,7 +1222,7 @@ class DenseGemmKernel: """ acc_shape = tiled_mma.partition_shape_C(mma_tiler[:2]) tCtAcc_fake = tiled_mma.make_fragment_C(acc_shape) - return sm100_utils.get_num_tmem_alloc_cols(tCtAcc_fake) + return utils.get_num_tmem_alloc_cols(tCtAcc_fake) def is_valid_dtypes( self, ab_dtype: Type[cutlass.Numeric], c_dtype: Type[cutlass.Numeric] @@ -1365,6 +1373,7 @@ class DenseGemmKernel: """ is_valid = True + # TODO: move to utils def check_contiguous_16B_alignment(dtype, is_mode0_major, tensor_shape): major_mode_idx = 0 if is_mode0_major else 1 num_major_elements = tensor_shape[major_mode_idx] @@ -1446,6 +1455,9 @@ class DenseGemmKernel: def create_tensors(l, m, n, k, a_major, b_major, c_major, ab_dtype, c_dtype): + import torch + import cutlass.torch as cutlass_torch + torch.manual_seed(1111) a_torch_cpu = cutlass_torch.matrix(l, m, k, a_major == "m", ab_dtype) @@ -1474,6 +1486,9 @@ def create_tensors(l, m, n, k, a_major, b_major, c_major, ab_dtype, c_dtype): def compare(a_torch_cpu, b_torch_cpu, c_torch_gpu, c_dtype, tolerance): + import torch + import cutlass.torch as cutlass_torch + # Copy gpu result back kernel_result = c_torch_gpu.cpu() @@ -1563,6 +1578,7 @@ def run( print(f"Iterations: {iterations}") print(f"Skip reference checking: {skip_ref_check}") print(f"Use cold L2: {'True' if use_cold_l2 else 'False'}") + import torch # Unpack parameters m, n, k, l = mnkl @@ -1600,6 +1616,8 @@ def run( compare(a_torch_cpu, b_torch_cpu, c_torch_gpu, c_dtype, tolerance) def generate_tensors(): + import cutlass.torch as cutlass_torch + a_tensor, _ = cutlass_torch.cute_tensor_like( a_torch_cpu, ab_dtype, is_dynamic_layout=True, assumed_align=16 ) diff --git a/examples/python/CuTeDSL/blackwell/epilogue/activation_custom_epilogue_dense_gemm.py b/examples/python/CuTeDSL/blackwell/epilogue/activation_custom_epilogue_dense_gemm.py new file mode 100644 index 00000000..5981ec7a --- /dev/null +++ b/examples/python/CuTeDSL/blackwell/epilogue/activation_custom_epilogue_dense_gemm.py @@ -0,0 +1,754 @@ +# 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. + +import traceback +import typing + +import cuda.bindings.driver as cuda + +# Required for pre-Python 3.12 instead of typing.override. +from typing_extensions import override +import torch + +import cutlass +import cutlass.cute.testing as testing +import cutlass.torch as cutlass_torch + +from common_dense_gemm_efc import DenseGemmEFC +from common_efc import ACTIVATION_FUNCTIONS + +""" +A high-performance persistent batched dense GEMM with activation functions in custom epilogue fusion +for the NVIDIA Blackwell SM100 architecture using CUTE DSL and Epilogue Fusion Configuration (EFC). + +This example demonstrates GEMMs with custom fused epilogues inspired by Ada FP8 GEMM epilogue +from https://github.com/NVIDIA/cutlass/blob/main/examples/58_ada_fp8_gemm/ada_fp8_gemm.cu : + Aux = ((alpha * scale_a * scale_b) * accumulator) + ((beta * scale_c) * source) + bias + D = activation(Aux) + +The scale factors (scale_a, scale_b, scale_c) default to 1.0 but can be customized via CLI: + +Tensor dimensions: +- Matrix A is MxKxL, L is batch dimension, A can be row-major("K") or column-major("M") +- Matrix B is NxKxL, L is batch dimension, B can be row-major("N") or column-major("K") +- Matrix C is MxNxL (read-only input, "source"), C can be row-major("N") or column-major("M") +- Matrix Aux is MxNxL (auxiliary output, pre-activation), same layout as C/D +- Matrix D is MxNxL (final output, post-activation), same layout as C +- alpha, beta are scalar scale factors +- scale_a, scale_b, scale_c are scalar scale factors for A, B, and C matrices +- bias is a scalar bias term + +This GEMM kernel supports the following features: + - Utilizes Tensor Memory Access (TMA) for efficient memory operations + - Utilizes Blackwell's tcgen05.mma for matrix multiply-accumulate (MMA) operations (including 2cta mma instructions) + - Implements TMA multicast with cluster to reduce L2 memory traffic + - Supports persistent tile scheduling to better overlap memory load/store with mma between tiles + - Supports warp specialization to avoid explicit pipelining between mainloop load and mma + - Uses Epilogue Fusion Configuration (EFC) to define custom epilogue operations with activation functions + +Supported activation functions: + - identity: f(x) = x + - relu: f(x) = max(0, x) + - leaky_relu: f(x) = max(0, x) + negative_slope * min(0, x) + - tanh: f(x) = tanh(x) + - sigmoid: f(x) = 1 / (1 + exp(-x)) + - silu: f(x) = x * sigmoid(x) + - hardswish: f(x) = x * relu6(x + 3) / 6 + - gelu: f(x) = 0.5 * x * (1 + erf(x / sqrt(2))) + +This GEMM works as follows: +1. DMA warp: Load A and B matrices from global memory (GMEM) to shared memory (SMEM) using TMA operations. +2. MMA warp: Perform matrix multiply-accumulate (MMA) operations using tcgen05.mma instruction. +3. EPILOGUE warp (defined via EFC): + - Load completed accumulator from tensor memory (TMEM) to registers (RMEM) using tcgen05.ld. + - Load C (source) matrix from global memory (GMEM) to shared memory (SMEM) using TMA, then to registers (RMEM). + - Compute Aux = (alpha * scale_a * scale_b) * accumulator + (beta * scale_c) * C + bias + - Compute D = activation(Aux) + - Type convert Aux and D matrices to output types. + - Store Aux and D matrices from registers (RMEM) to shared memory (SMEM) to global memory (GMEM) with TMA operations + +SM100 tcgen05.mma instructions operate as follows: +- Read matrix A from SMEM +- Read matrix B from SMEM +- Write accumulator to TMEM +The accumulator in TMEM must then be loaded to registers before writing back to GMEM. + +Example usage: + +.. code-block:: bash + + python activation_custom_epilogue_dense_gemm.py \ + --activation relu \ + --ab_dtype Float16 --c_dtype Float16 --aux_dtype Float16 --d_dtype Float16 \ + --acc_dtype Float32 --epi_dtype Float32 \ + --mma_tiler_mn 128,128 --cluster_shape_mn 2,1 \ + --mnkl 8192,8192,8192,1 \ + --use_2cta_instrs --alpha 1.0 --beta 1.0 --bias 0.0 \ + --scale_a 1.0 --scale_b 1.0 --scale_c 1.0 + +To collect performance with NCU profiler: + +.. code-block:: bash + + ncu python activation_custom_epilogue_dense_gemm.py \ + --activation gelu \ + --ab_dtype Float16 --c_dtype Float16 --aux_dtype Float16 --d_dtype Float16 \ + --acc_dtype Float32 --epi_dtype Float32 \ + --mma_tiler_mn 128,128 --cluster_shape_mn 2,1 \ + --mnkl 8192,8192,8192,1 \ + --use_2cta_instrs --alpha 1.0 --beta 1.0 --bias 0.0 \ + --scale_a 1.0 --scale_b 1.0 --scale_c 1.0 \ + --warmup_iterations 1 --iterations 10 --skip_ref_check + +Constraints: +* Supported input data types: fp16, bf16, tf32, int8, uint8, fp8 (e4m3fn, e5m2) +* A/B tensors must have the same data type +* C/D/Aux tensors must have the same major order +* MMA tiler M must be 64/128 (use_2cta_instrs=False) or 128/256 (use_2cta_instrs=True) +* MMA tiler N must be 32-256, step 32 +* Cluster shape M/N must be positive and power of 2, total cluster size <= 16 +* Cluster shape M must be multiple of 2 if use_2cta_instrs=True +* The contiguous dimension of all tensors must be at least 16 bytes aligned, + i.e, number of elements is a multiple of 4, 8, and 16 for TFloat32, + Float16/BFloat16, and Int8/Uint8/Float8, respectively. +* OOB tiles are not allowed when TMA store is disabled +""" + + +class DenseGemmActivation(DenseGemmEFC): + """Implements batched GEMM with activation function in epilogue using EFC. + + This class extends DenseGemmEFC to provide a fused epilogue inspired by + Ada FP8 GEMM that: + - Reads from input tensor C (source) + - Writes to output tensors Aux (auxiliary, pre-activation) and D (final, post-activation) + - Performs: Aux = alpha * accumulator + beta * C + bias + D = activation(Aux) + + The class provides CLI argument parsing and tensor creation for the + specific epilogue configuration with C, Aux, D tensors and alpha, + beta, bias scalar parameters. + """ + + def __init__( + self, + acc_dtype, + epi_dtype, + use_2cta_instrs, + mma_tiler_mn, + cluster_shape_mn, + epilogue_fn, + activation_name, + ): + """Initialize the GEMM with activation epilogue. + + :param acc_dtype: Accumulator data type + :param epi_dtype: Epilogue computation data type + :param use_2cta_instrs: Whether to use 2-CTA MMA instructions + :param mma_tiler_mn: MMA tile shape (M, N) + :param cluster_shape_mn: Cluster shape (M, N) + :param epilogue_fn: Epilogue function to use + :param activation_name: Name of the activation function + """ + super().__init__( + acc_dtype, + epi_dtype, + use_2cta_instrs, + mma_tiler_mn, + cluster_shape_mn, + epilogue_fn, + ) + self.activation_name = activation_name + + class CLIParser(DenseGemmEFC.CLIParser): + @override + def more_parsing(self): + self.parser.add_argument( + "--activation", + type=str, + default="relu", + choices=ACTIVATION_FUNCTIONS, + help="Activation function to use in epilogue", + ) + self.parser.add_argument( + "--alpha", + type=float, + default=1.0, + help="alpha scale factor for accumulator", + ) + self.parser.add_argument( + "--beta", type=float, default=1.0, help="beta scale factor for source" + ) + self.parser.add_argument( + "--bias", type=float, default=0.0, help="bias term to add" + ) + self.parser.add_argument( + "--scale_a", type=float, default=1.0, help="scale factor for matrix A" + ) + self.parser.add_argument( + "--scale_b", type=float, default=1.0, help="scale factor for matrix B" + ) + self.parser.add_argument( + "--scale_c", type=float, default=1.0, help="scale factor for source C" + ) + self.parser.add_argument( + "--c_dtype", + type=cutlass.dtype, + default=cutlass.Float32, + help="C tensor dtype", + ) + self.parser.add_argument( + "--aux_dtype", + type=cutlass.dtype, + default=cutlass.Float32, + help="Aux tensor dtype", + ) + self.parser.add_argument( + "--d_dtype", + type=cutlass.dtype, + default=cutlass.Float32, + help="D tensor dtype", + ) + self.parser.add_argument( + "--leaky_relu_alpha", + type=float, + default=0.01, + help="negative slope for leaky_relu", + ) + + @override + def create_arguments( + self, + l, + m, + n, + k, + a_major, + b_major, + cd_major, + ab_dtype, + # For the supplemental tensors. + c_dtype, + aux_dtype, + d_dtype, + ): + """Create arguments for GEMM operations with epilogue tensors. + + Creates tensors for A, B (from parent class) and epilogue-specific + tensors C, Aux, D with appropriate data types and layouts. + + :return: Tuple of (a_tensor, b_tensor, a_torch_cpu, b_torch_cpu, + c_tensor, c_torch_cpu, c_torch_gpu, + aux_tensor, aux_torch_cpu, aux_torch_gpu, + d_tensor, d_torch_cpu, d_torch_gpu) + """ + # Get standard arguments from parent class + std_args = super().create_arguments( + l, m, n, k, a_major, b_major, cd_major, ab_dtype + ) + + # Create C tensor (source for epilogue) + c_torch_cpu = cutlass_torch.matrix(l, m, n, cd_major == "m", c_dtype) + c_tensor, c_torch_gpu = cutlass_torch.cute_tensor_like( + c_torch_cpu, c_dtype, is_dynamic_layout=True, assumed_align=16 + ) + + # Create Aux tensor (auxiliary/pre-activation output) + aux_torch_cpu = cutlass_torch.matrix(l, m, n, cd_major == "m", aux_dtype) + aux_tensor, aux_torch_gpu = cutlass_torch.cute_tensor_like( + aux_torch_cpu, aux_dtype, is_dynamic_layout=True, assumed_align=16 + ) + + # Create D tensor (final/post-activation output) + d_torch_cpu = cutlass_torch.matrix(l, m, n, cd_major == "m", d_dtype) + d_tensor, d_torch_gpu = cutlass_torch.cute_tensor_like( + d_torch_cpu, d_dtype, is_dynamic_layout=True, assumed_align=16 + ) + + return ( + *std_args, + c_tensor, + c_torch_cpu, + c_torch_gpu, + aux_tensor, + aux_torch_cpu, + aux_torch_gpu, + d_tensor, + d_torch_cpu, + d_torch_gpu, + ) + + def compare( + self, + a_torch_cpu, + b_torch_cpu, + epi_dtype, + tolerance, + # For the tensor check. + c_torch_gpu, + aux_torch_gpu, + d_torch_gpu, + # The EFC epilogue arguments. + c_torch_cpu, + aux_torch_cpu, + alpha, + beta, + bias, + scale_a, + scale_b, + scale_c, + d_torch_cpu, + leaky_relu_alpha=0.01, + ): + """Compare GPU results against CPU reference implementation. + + :param a_torch_cpu: Input tensor A on CPU + :param b_torch_cpu: Input tensor B on CPU + :param epi_dtype: Epilogue data type + :param tolerance: Comparison tolerance + :param c_torch_gpu: GPU result for C + :param aux_torch_gpu: GPU result for Aux + :param d_torch_gpu: GPU result for D + :param c_torch_cpu: CPU reference for C + :param aux_torch_cpu: CPU reference for Aux + :param alpha: Alpha scale factor + :param beta: Beta scale factor + :param bias: Bias term + :param scale_a: Scale factor for matrix A + :param scale_b: Scale factor for matrix B + :param scale_c: Scale factor for source C + :param d_torch_cpu: CPU reference for D + :param leaky_relu_alpha: Negative slope for leaky_relu + """ + # Compute reference result + self.evaluate_on_cpu( + a_torch_cpu, + b_torch_cpu, + epi_dtype, + c_torch_cpu, + aux_torch_cpu, + alpha, + beta, + bias, + scale_a, + scale_b, + scale_c, + d_torch_cpu, + leaky_relu_alpha, + ) + # Assert close results for output tensors + torch.testing.assert_close( + aux_torch_gpu.cpu(), aux_torch_cpu, atol=tolerance, rtol=1e-03 + ) + torch.testing.assert_close( + d_torch_gpu.cpu(), d_torch_cpu, atol=tolerance, rtol=1e-03 + ) + # Assert that the read tensor has not been changed + torch.testing.assert_close( + c_torch_gpu.cpu(), c_torch_cpu, atol=tolerance, rtol=1e-03 + ) + + @staticmethod + def format_as_cli_args( + mnkl: typing.Tuple[int, int, int, int], + ab_dtype: typing.Type[cutlass.Numeric], + acc_dtype: typing.Type[cutlass.Numeric], + epi_dtype: typing.Type[cutlass.Numeric], + a_major: str, + b_major: str, + cd_major: str, + mma_tiler_mn: typing.Tuple[int, int], + cluster_shape_mn: typing.Tuple[int, int], + use_2cta_instrs: bool, + c_dtype: typing.Type[cutlass.Numeric], + aux_dtype: typing.Type[cutlass.Numeric], + d_dtype: typing.Type[cutlass.Numeric], + alpha: float, + beta: float, + bias: float, + scale_a: float, + scale_b: float, + scale_c: float, + activation: str, + leaky_relu_alpha: float, + tolerance: float, + ) -> str: + """Format test parameters as CLI arguments for activation_custom_epilogue_dense_gemm.py + + Formats all test parameters into a CLI command that can be directly + copy-pasted to reproduce the test case. Includes base parameters from + DenseGemmEFC and epilogue-specific parameters (c_dtype, aux_dtype, d_dtype, + alpha, beta, bias, activation). + + :return: Formatted CLI command string + """ + # Get base command from parent class + base_cmd = DenseGemmEFC.format_as_cli_args( + "activation_custom_epilogue_dense_gemm.py", + mnkl, + ab_dtype, + acc_dtype, + epi_dtype, + a_major, + b_major, + cd_major, + mma_tiler_mn, + cluster_shape_mn, + use_2cta_instrs, + tolerance, + ) + + # Add epilogue-specific arguments + epilogue_args = ( + f" --activation {activation}" + f" --c_dtype {c_dtype.__name__}" + f" --aux_dtype {aux_dtype.__name__}" + f" --d_dtype {d_dtype.__name__}" + f" --alpha {alpha}" + f" --beta {beta}" + f" --bias {bias}" + f" --scale_a {scale_a}" + f" --scale_b {scale_b}" + f" --scale_c {scale_c}" + f" --leaky_relu_alpha {leaky_relu_alpha}" + ) + + return base_cmd + epilogue_args + + +def create_epilogue_function(activation_name: str): + """Create an epilogue function with the specified activation. + + :param activation_name: Name of the activation function to use + :return: Epilogue function + """ + # Validate activation name + if activation_name not in ACTIVATION_FUNCTIONS: + raise ValueError(f"Unsupported activation: {activation_name}") + + def epilogue( + efc_config, + C, + Aux, + alpha, + beta, + bias, + scale_a, + scale_b, + scale_c, + D, + leaky_relu_alpha, + ): + # Aux = ((alpha * scale_a * scale_b) * accumulator) + ((beta * scale_c) * source) + bias + # Following Ada FP8 GEMM epilogue pattern + aux_val = ( + (alpha * scale_a * scale_b) * efc_config.accum() + + (beta * scale_c) * C.load() + + bias + ) + Aux.store(aux_val) + # D = activation(Aux) + activation_fn = getattr(efc_config, activation_name) + # leaky_relu needs an extra parameter, others don't + if activation_name == "leaky_relu": + D.store(activation_fn(aux_val, leaky_relu_alpha)) + else: + D.store(activation_fn(aux_val)) + + return epilogue + + +def run( + mnkl: typing.Tuple[int, int, int, int], + ab_dtype: typing.Type[cutlass.Numeric], + acc_dtype: typing.Type[cutlass.Numeric], + epi_dtype: typing.Type[cutlass.Numeric], + a_major: str, + b_major: str, + cd_major: str, + c_dtype: typing.Type[cutlass.Numeric], + aux_dtype: typing.Type[cutlass.Numeric], + d_dtype: typing.Type[cutlass.Numeric], + alpha: float, + beta: float, + bias: float, + scale_a: float, + scale_b: float, + scale_c: float, + activation: str, + leaky_relu_alpha: float, + mma_tiler_mn: typing.Tuple[int, int], + cluster_shape_mn: typing.Tuple[int, int], + use_2cta_instrs: bool, + tolerance: float, + warmup_iterations: int = 3, + iterations: int = 100, + skip_ref_check: bool = False, +): + """Run GEMM with activation function in epilogue. + + :param mnkl: Tuple of (M, N, K, L) dimensions + :param ab_dtype: Data type for A and B tensors + :param acc_dtype: Accumulator data type + :param epi_dtype: Epilogue computation data type + :param a_major: Major dimension for A ("m" or "k") + :param b_major: Major dimension for B ("n" or "k") + :param cd_major: Major dimension for C/D/Aux ("m" or "n") + :param c_dtype: Data type for C tensor + :param aux_dtype: Data type for Aux tensor + :param d_dtype: Data type for D tensor + :param alpha: Alpha scale factor + :param beta: Beta scale factor + :param bias: Bias term + :param scale_a: Scale factor for matrix A + :param scale_b: Scale factor for matrix B + :param scale_c: Scale factor for source C + :param activation: Activation function name + :param leaky_relu_alpha: Negative slope for leaky_relu + :param mma_tiler_mn: MMA tile shape (M, N) + :param cluster_shape_mn: Cluster shape (M, N) + :param use_2cta_instrs: Whether to use 2-CTA MMA instructions + :param tolerance: Comparison tolerance + :param warmup_iterations: Number of warmup iterations + :param iterations: Number of benchmark iterations + :param skip_ref_check: Whether to skip reference check + """ + print("Running Blackwell Persistent Dense GEMM test with:") + print(f"mnkl: {mnkl}") + print(f"AB dtype: {ab_dtype}, Acc dtype: {acc_dtype}, Epi dtype: {epi_dtype}") + print( + f"Matrix majors - A: {a_major}, B: {b_major}, loaded: {cd_major}, stored: {cd_major}" + ) + print(f"Mma Tiler (M, N): {mma_tiler_mn}, Cluster Shape (M, N): {cluster_shape_mn}") + print(f"2CTA MMA instructions: {'True' if use_2cta_instrs else 'False'}") + print(f"Tolerance: {tolerance}") + print(f"Warmup iterations: {warmup_iterations}") + print(f"Iterations: {iterations}") + print(f"Skip reference checking: {skip_ref_check}") + print("Epilogue:") + print(f"\t{c_dtype = !s}, {aux_dtype = !s}, {d_dtype = !s}") + print(f"\t{alpha = }, {beta = }, {bias = }") + print(f"\t{scale_a = }, {scale_b = }, {scale_c = }") + print(f"\t{activation = !s}") + if activation == "leaky_relu": + print(f"\t{leaky_relu_alpha = }") + + # Unpack parameters + m, n, k, l = mnkl + + if not torch.cuda.is_available(): + raise RuntimeError("GPU is required to run this example!") + + # Get current CUDA stream from PyTorch + torch_stream = torch.cuda.current_stream() + # Get the raw stream pointer as a CUstream + current_stream = cuda.CUstream(torch_stream.cuda_stream) + + # Create the epilogue function with the specified activation + epilogue_fn = create_epilogue_function(activation) + + # Build GEMM object with EFC configuration + gemm = DenseGemmActivation( + acc_dtype, + epi_dtype, + use_2cta_instrs, + mma_tiler_mn, + cluster_shape_mn, + epilogue_fn, + activation, + ) + + ( + a_tensor, + b_tensor, + a_torch_cpu, + b_torch_cpu, + # The supplemental tensors. + c_tensor, + c_torch_cpu, + c_torch_gpu, + aux_tensor, + aux_torch_cpu, + aux_torch_gpu, + d_tensor, + d_torch_cpu, + d_torch_gpu, + ) = gemm.create_arguments( + l, + m, + n, + k, + a_major, + b_major, + cd_major, + ab_dtype, + # For the supplemental tensors. + c_dtype, + aux_dtype, + d_dtype, + ) + + # Check if the configuration can be implemented. Raise a ValueError + # otherwise. + gemm.check_implementable(a_tensor, b_tensor, d_tensor) + + max_active_clusters = cutlass.utils.HardwareInfo().get_max_active_clusters( + cluster_shape_mn[0] * cluster_shape_mn[1] + ) + + compiled_gemm = gemm.compile( + a_tensor, + b_tensor, + max_active_clusters, + current_stream, + # Here are the supplemental arguments in the same order as for the + # epilogue configuration function. + c_tensor, + aux_tensor, + alpha, + beta, + bias, + scale_a, + scale_b, + scale_c, + d_tensor, + leaky_relu_alpha, + ) + + compiled_gemm( + a_tensor, + b_tensor, + current_stream, + # Here are the supplemental arguments in the same order as for the + # epilogue configuration function. + c_tensor, + aux_tensor, + alpha, + beta, + bias, + scale_a, + scale_b, + scale_c, + d_tensor, + leaky_relu_alpha, + ) + + # TODO: unify with modern way to do benchmarking. + exec_time = testing.benchmark( + compiled_gemm, + kernel_arguments=testing.JitArguments( + a_tensor, + b_tensor, + current_stream, + # Here are the supplemental arguments in the same order as for the + # epilogue configuration function. + c_tensor, + aux_tensor, + alpha, + beta, + bias, + scale_a, + scale_b, + scale_c, + d_tensor, + leaky_relu_alpha, + ), + stream=current_stream, + warmup_iterations=warmup_iterations, + iterations=iterations, + ) + + print(f"Execution time: {exec_time} us") + + # Compute reference result + if not skip_ref_check: + print("Checking results against CPU reference...") + gemm.compare( + # The usual arguments. + a_torch_cpu, + b_torch_cpu, + epi_dtype, + tolerance, + # For the tensor check. + c_torch_gpu, + aux_torch_gpu, + d_torch_gpu, + # The EFC epilogue arguments. + c_torch_cpu, + aux_torch_cpu, + alpha, + beta, + bias, + scale_a, + scale_b, + scale_c, + d_torch_cpu, + leaky_relu_alpha, + ) + print("Results match CPU reference!") + + +if __name__ == "__main__": + args = DenseGemmActivation.CLIParser().parse() + + try: + run( + args.mnkl, + args.ab_dtype, + args.acc_dtype, + args.epi_dtype, + args.a_major, + args.b_major, + args.cd_major, + args.c_dtype, + args.aux_dtype, + args.d_dtype, + args.alpha, + args.beta, + args.bias, + args.scale_a, + args.scale_b, + args.scale_c, + args.activation, + args.leaky_relu_alpha, + args.mma_tiler_mn, + args.cluster_shape_mn, + args.use_2cta_instrs, + args.tolerance, + args.warmup_iterations, + args.iterations, + args.skip_ref_check, + ) + print("\n" + "=" * 80) + print( + f"PASS - {args.activation.upper()} activation test completed successfully!" + ) + print("=" * 80 + "\n") + except Exception as exc: + traceback.print_exception(exc) + raise diff --git a/examples/python/CuTeDSL/blackwell/epilogue/common_dense_gemm_efc.py b/examples/python/CuTeDSL/blackwell/epilogue/common_dense_gemm_efc.py new file mode 100644 index 00000000..0a4fa2d7 --- /dev/null +++ b/examples/python/CuTeDSL/blackwell/epilogue/common_dense_gemm_efc.py @@ -0,0 +1,2011 @@ +# 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. + +import argparse +import logging +import types +import typing +from typing import Tuple, Type, Union + +import cuda.bindings.driver as cuda +import torch + +import cutlass +import cutlass.cute as cute +import cutlass.pipeline as pipeline +import cutlass.torch as cutlass_torch +import cutlass.utils as utils +import cutlass.utils.blackwell_helpers as sm100_utils +from cutlass.cute.nvgpu import cpasync, tcgen05 + +import common_efc +from common_efc import log + +""" +Common base infrastructure for high-performance persistent batched dense GEMM with custom epilogue fusion +for the NVIDIA Blackwell SM100 architecture using CUTE DSL and Epilogue Fusion Configuration (EFC). + +This module provides the DenseGemmEFC base class that implements the core GEMM functionality with +support for custom epilogue operations. Subclasses define specific epilogue configurations by +providing an epilogue function that operates on the accumulator and supplemental tensors. + +Key Features: + - Utilizes Tensor Memory Access (TMA) for efficient memory operations + - Utilizes Blackwell's tcgen05.mma for matrix multiply-accumulate (MMA) operations (including 2cta mma instructions) + - Implements TMA multicast with cluster to reduce L2 memory traffic + - Supports persistent tile scheduling to better overlap memory load/store with mma between tiles + - Supports warp specialization to avoid explicit pipelining between mainloop load and mma + - Uses Epilogue Fusion Configuration (EFC) to define custom epilogue operations + +GEMM Execution Flow: +1. DMA warp: Load A and B matrices from global memory (GMEM) to shared memory (SMEM) using TMA operations. +2. MMA warp: Perform matrix multiply-accumulate (MMA) operations using tcgen05.mma instruction. +3. EPILOGUE warp (customizable via EFC): + - Load completed accumulator from tensor memory (TMEM) to registers (RMEM) using tcgen05.ld. + - Load supplemental input tensors from GMEM to SMEM using TMA, then to RMEM. + - Execute custom epilogue function (defined by subclass) that: + * Accesses the accumulator via efc_config.accum() + * Reads from supplemental input tensors via tensor.load() + * Writes to supplemental output tensors via tensor.store() + * Can apply arbitrary element-wise operations, scaling, and fusion + - Store result tensors from RMEM to SMEM to GMEM with TMA operations. + +SM100 tcgen05.mma instructions operate as follows: +- Read matrix A from SMEM +- Read matrix B from SMEM +- Write accumulator to TMEM +The accumulator in TMEM must then be loaded to registers before writing back to GMEM. + +Base Tensor Dimensions: +- Matrix A is MxKxL, L is batch dimension, A can be row-major("K") or column-major("M") +- Matrix B is NxKxL, L is batch dimension, B can be row-major("N") or column-major("K") +- Supplemental tensors are MxNxL with layout matching the epilogue configuration + +Common Constraints: +* Supported input data types: fp16, bf16, tf32, int8, uint8, fp8 (e4m3fn, e5m2) +* A/B tensors must have the same data type +* MMA tiler M must be 64/128 (use_2cta_instrs=False) or 128/256 (use_2cta_instrs=True) +* MMA tiler N must be 32-256, step 32 +* Cluster shape M/N must be positive and power of 2, total cluster size <= 16 +* Cluster shape M must be multiple of 2 if use_2cta_instrs=True +* The contiguous dimension of all tensors must be at least 16 bytes aligned, + i.e, number of elements is a multiple of 4, 8, and 16 for TFloat32, + Float16/BFloat16, and Int8/Uint8/Float8, respectively. +* OOB tiles are not allowed when TMA store is disabled + +Subclass Examples: +- custom_epilogue_dense_gemm.py: Custom fused epilogue with multiple read/write tensors +- synthetic_custom_epilogue_dense_gemm.py: Synthetic epilogue for testing with configurable tensor counts +""" + + +class DenseGemmEFC: + """Base class for batched GEMM with custom epilogue fusion using EFC. + + This class provides the core infrastructure for persistent batched GEMM operations + with customizable epilogue fusion. Subclasses define specific epilogue behaviors + by providing an epilogue configuration function that describes operations on the + accumulator and supplemental tensors. + + The class handles: + - GEMM mainloop (A * B computation) + - TMA-based memory operations + - Warp specialization + - Persistent tile scheduling + - EFC (Epilogue Fusion Configuration) integration + - CLI argument parsing (extensible via CLIParser.more_parsing()) + - Tensor creation and validation + + :param acc_dtype: Data type for accumulation during computation + :type acc_dtype: type[cutlass.Numeric] + :param epi_dtype: Data type for epilogue operation + :type epi_dtype: type[cutlass.Numeric] + :param use_2cta_instrs: Whether to use CTA group 2 for advanced thread cooperation + :type use_2cta_instrs: bool + :param mma_tiler_mn: Shape of the Matrix Multiply-Accumulate (MMA) tile (M,N) + :type mma_tiler_mn: Tuple[int, int] + :param cluster_shape_mn: Cluster dimensions (M,N) for parallel processing + :type cluster_shape_mn: Tuple[int, int] + :param epilogue_function_configuration: Function defining the epilogue behavior via EFC + :type epilogue_function_configuration: Callable + + :note: Supported A/B data types: + - TFloat32 + - Float16/BFloat16 + - Int8/Uint8 + - Float8E4M3FN/Float8E5M2 + (A and B must have the same data type) + + :note: Supported accumulator data types: + - Float32 (for all floating point A/B data types) + - Float16 (only for fp16 and fp8 A/B data types) + - Int32 (only for uint8/int8 A/B data types) + + :note: Supported supplemental tensor data types (epilogue-dependent): + - Float32 (for float32 and int32 accumulator data types) + - Int32 (for float32 and int32 accumulator data types) + - Float16/BFloat16 (for fp16 and fp8 accumulator data types) + - Int8/Uint8 (for uint8/int8 accumulator data types) + - Float8E4M3FN/Float8E5M2 (for float32 accumulator data types) + + :note: Constraints: + - MMA tiler M must be 64/128 (use_2cta_instrs=False) or 128/256 (use_2cta_instrs=True) + - MMA tiler N must be 32-256, step 32 + - Cluster shape M must be multiple of 2 if use_2cta_instrs=True + - Cluster shape M/N must be positive and power of 2, total cluster size <= 16 + + Example: + >>> def my_epilogue(efc_config, alpha, beta, output_tensor, input_tensor): + ... result = efc_config.accum() * alpha + input_tensor.load() * beta + ... output_tensor.store(result) + ... + >>> gemm = DenseGemmEFC( + ... acc_dtype=cutlass.Float32, + ... epi_dtype=cutlass.Float32, + ... use_2cta_instrs=True, + ... mma_tiler_mn=(128, 128), + ... cluster_shape_mn=(2, 2), + ... epilogue_function_configuration=my_epilogue + ... ) + """ + + def __init__( + self, + acc_dtype: Type[cutlass.Numeric], + epi_dtype: Type[cutlass.Numeric], + use_2cta_instrs: bool, + mma_tiler_mn: Tuple[int, int], + cluster_shape_mn: Tuple[int, int], + epilogue_function_configuration: typing.Callable, + ): + """Initializes the configuration for a Blackwell dense GEMM kernel with EFC. + + This configuration includes several key aspects: + + 1. MMA Instruction Settings (tcgen05): + - acc_dtype: Data types for MMA accumulator. + - mma_tiler_mn: The (M, N) shape of the MMA instruction tiler. + - use_2cta_instrs: Boolean indicating if the tcgen05 MMA variant + with cta_group=2 should be used. + + 2. Cluster Shape: + - cluster_shape_mn: The (ClusterM, ClusterN) shape of the CTA cluster. + + 3. Epilogue Configuration: + - epilogue_function_configuration: Defines custom epilogue behavior + that operates on accumulator and supplemental tensors. + + :param acc_dtype: Data type of the accumulator. + :type acc_dtype: type[cutlass.Numeric] + :param epi_dtype: Data type of the epilogue. + :type epi_dtype: type[cutlass.Numeric] + :param use_2cta_instrs: Boolean, True to use cta_group=2 MMA variant. + :type use_2cta_instrs: bool + :param mma_tiler_mn: Tuple (M, N) shape of the MMA instruction. + :type mma_tiler_mn: Tuple[int, int] + :param cluster_shape_mn: Tuple (ClusterM, ClusterN) shape of the cluster. + :type cluster_shape_mn: Tuple[int, int] + :param epilogue_function_configuration: Function defining epilogue behavior via EFC. + :type epilogue_function_configuration: Callable + """ + + self.acc_dtype: Type[cutlass.Numeric] = acc_dtype + self.epi_dtype: Type[cutlass.Numeric] = epi_dtype + self.use_2cta_instrs = use_2cta_instrs + self.cluster_shape_mn = cluster_shape_mn + # K dimension is deferred in _setup_attributes + self.mma_tiler_mn = mma_tiler_mn + self.mma_tiler = (*mma_tiler_mn, 1) + self.arch = "sm_100" + + self.c_dtype = self.epi_dtype + + self.cta_group = ( + tcgen05.CtaGroup.TWO if self.use_2cta_instrs else tcgen05.CtaGroup.ONE + ) + + self.occupancy = 1 + # Set specialized warp ids: + + # The warps responsible for computing the epilogue function and storing + # the results. + self.epilogue_warp_id = (0, 1, 2, 3) + # The warp responsible for computing the matrix multiplication. + self.mma_warp_id = 4 + # The warp responsible for loading the tensors A & B to feed the MMA. + self.tma_warp_id = 5 + # The warp responsible for loading the auxiliary tensors used in the epilogue. + self.epilogue_load_warp_id = 6 + self.threads_per_cta = 32 * len( + ( + self.mma_warp_id, + self.tma_warp_id, + *self.epilogue_warp_id, + self.epilogue_load_warp_id, + ) + ) + # Barrier ids for cta sync, epilogue sync and tmem ptr sync. + self.cta_sync_bar_id = 1 + self.epilogue_sync_bar_id = 2 + self.tmem_alloc_sync_bar_id = 3 + # Amount of available shared memory. + self.smem_capacity = utils.get_smem_capacity_in_bytes(self.arch) + + # Setup the EFC from the given function representing the epilogue + # configuration. + self.efc = common_efc.EFC(self, epilogue_function_configuration) + + def _create_tiled_mma(self): + """Make a tiled MMA atom with given data type, leading dimension, CTA + group and MMA tile shape. Use SMEM operand source for A.""" + return utils.sm100.make_trivial_tiled_mma( + self.a_dtype, + self.a_major_mode, + self.b_major_mode, + self.acc_dtype, + self.cta_group, + self.mma_tiler[:2], + ) + + def _setup_attributes(self): + """Set up configurations that are dependent on GEMM inputs + + This method configures various attributes based on the input tensor properties + (data types, leading dimensions) and kernel settings: + - Configuring tiled MMA + - Computing MMA/cluster/tile shapes + - Computing cluster layout + - Computing multicast CTAs for A/B + - Computing epilogue subtile + - Setting up A/B/C/D stage counts in shared memory + - Computing A/B/C/D shared memory layout + - Computing tensor memory allocation columns + """ + # Get the right tiled MMA. + self._tiled_mma = self._create_tiled_mma() + log(f"{self._tiled_mma = !s}") + + # Compute mma/cluster/tile shapes + mma_inst_shape_k = cute.size(self._tiled_mma.shape_mnk, mode=[2]) + mma_inst_tile_k = 4 + # Extend mma_tiler with k-dimension (MMA_M, MMA_N, MMA_K) + self.mma_tiler = ( + self.mma_tiler[0], + self.mma_tiler[1], + mma_inst_shape_k * mma_inst_tile_k, + ) + log(f"{self.mma_tiler = !s}") + # CTA tiler with the 2CTA instruction correction. + self.cta_tile_shape_mnk = ( + self.mma_tiler[0] // cute.size(self._tiled_mma.thr_id.shape), + self.mma_tiler[1], + self.mma_tiler[2], + ) + log(f"{self.cta_tile_shape_mnk = !s}") + # Compute cluster layout, V for the 2CTA instructions. + self.cluster_layout_vmnk = cute.tiled_divide( + cute.make_layout((*self.cluster_shape_mn, 1)), + (self._tiled_mma.thr_id.shape,), + ) + log(f"{cute.make_layout((*self.cluster_shape_mn, 1)) = !s}") + log(f"{self.cluster_layout_vmnk = !s}") + # Compute number of multicast CTAs for A/B + self.num_mcast_ctas_a = cute.size(self.cluster_layout_vmnk.shape[2]) + self.num_mcast_ctas_b = cute.size(self.cluster_layout_vmnk.shape[1]) + self.is_a_mcast = self.num_mcast_ctas_a > 1 + self.is_b_mcast = self.num_mcast_ctas_b > 1 + log(f"{self.num_mcast_ctas_a = }, {self.num_mcast_ctas_b = }") + log(f"{self.is_a_mcast = }, {self.is_b_mcast = }") + + # Compute epilogue (EPI_TILE_M, EPI_TILE_N) subtile of cta_tile_shape_mnk + # according to some heuristics. + self.epi_tile = sm100_utils.compute_epilogue_tile_shape( + self.cta_tile_shape_mnk, + self.use_2cta_instrs, + layout_d=self.d_layout, + elem_ty_d=self.d_dtype, + layout_c=self.c_layout, + elem_ty_c=self.c_dtype, + ) + log(f"{self.epi_tile = !s}") + + # Setup A/B/C/D pipeline stage count in shared memory and ACC stage + # count in tensor memory. + self.compute_stages() + log(f"{self.num_acc_stage = }, {self.num_ab_stage = }, {self.num_c_stage = }") + # Compute A/B shared memory layout + self.a_smem_layout_staged = sm100_utils.make_smem_layout_a( + self._tiled_mma, + self.mma_tiler, + self.a_dtype, + self.num_ab_stage, + ) + log(f"{self.a_smem_layout_staged = !s}") + self.b_smem_layout_staged = sm100_utils.make_smem_layout_b( + self._tiled_mma, + self.mma_tiler, + self.b_dtype, + self.num_ab_stage, + ) + log(f"{self.b_smem_layout_staged = !s}") + # Get the smem_layout for the tensors used in the EFC. + self.efc.jit.smem_layout() + + # Compute the number of tensor memory allocation columns + self.compute_num_tmem_alloc_cols() + + @cute.jit + def __call__( + self, + a: cute.Tensor, + b: cute.Tensor, + max_active_clusters: cutlass.Constexpr, + stream: cuda.CUstream, + supplemental_parameters: Tuple, + ): + """Execute the GEMM operation in steps: + - Setup static attributes before smem/grid/tma computation + - Setup TMA load/store atoms and tensors + - Compute grid size with regard to hardware constraints + - Define shared storage for kernel + - Launch the kernel synchronously + + :param a: Input tensor A + :type a: cute.Tensor + :param b: Input tensor B + :type b: cute.Tensor + :param max_active_clusters: Maximum number of active clusters + :type max_active_clusters: cutlass.Constexpr + :param stream: CUDA stream for asynchronous execution + :type stream: cuda.CUstream + :param supplemental_parameters: Tuple or None used to pass variadic values + :type supplemental_parameters: Tuple + :raises TypeError: If input data types are incompatible with the MMA instruction. + :raises AssertionError: If OOB (Out-Of-Bounds) tiles are present when TMA store is disabled. + """ + # Process the variadic parameters. + self.efc.jit.unpack_parameters(supplemental_parameters) + + # Setup static attributes before smem/grid/tma computation + self.a_dtype: Type[cutlass.Numeric] = a.element_type + self.b_dtype: Type[cutlass.Numeric] = b.element_type + self.a_major_mode = utils.LayoutEnum.from_tensor(a).mma_major_mode() + self.b_major_mode = utils.LayoutEnum.from_tensor(b).mma_major_mode() + + # Gather all the auxiliary tensor element data types. + self.efc.jit.record_tensor_dtypes() + + # There is no D tensor to be used as a returned tensor. In the + # following, D is used more like of a "store" concept. So use the + # written tensor with the biggest element_type to set up all the tiling + # heuristics and epilogue store pipeline. + self.d_name_bigger = self.efc.jit.written_tensor_name_with_bigger_element_type() + d = self.efc.jit.parameter[self.d_name_bigger] + self.d_dtype: Type[cutlass.Numeric] = d.element_type + self.d_layout = utils.LayoutEnum.from_tensor(d) + log(f"d{self.d_name_bigger} = {d!s}") + + # C is the read tensor with the biggest element_type, if any, used by + # some heuristics for tiling. + self.c_dtype = None + self.c_layout = None + self.c_name_bigger = self.efc.jit.read_tensor_name_with_bigger_element_type() + if cutlass.const_expr(self.c_name_bigger): + c = self.efc.jit.parameter[self.c_name_bigger] + log(f"{self.c_name_bigger = } -> {c = !s}") + self.c_dtype = c.element_type + self.c_layout = utils.LayoutEnum.from_tensor(c) + + # Check if input data types are compatible with MMA instruction + if cutlass.const_expr(self.a_dtype != self.b_dtype): + raise TypeError(f"Type must match: {self.a_dtype} != {self.b_dtype}") + + # Setup attributes that depend on gemm inputs + self._setup_attributes() + + atom_thr_size = cute.size(self._tiled_mma.thr_id.shape) + + # Setup TMA load for A + a_op = utils.sm100.cluster_shape_to_tma_atom_A( + self.cluster_shape_mn, self._tiled_mma.thr_id + ) + log(f"{a_op = !s}") + # Get read of the pipeline dimension. + a_smem_layout = cute.slice_(self.a_smem_layout_staged, (None, None, None, 0)) + log(f"{a_smem_layout = !s}") + tma_atom_a, tma_tensor_a = cute.nvgpu.make_tiled_tma_atom_A( + a_op, + a, + a_smem_layout, + self.mma_tiler, + self._tiled_mma, + self.cluster_layout_vmnk.shape, + internal_type=( + cutlass.TFloat32 if a.element_type is cutlass.Float32 else None + ), + ) + log(f"{tma_atom_a = !s}") + log(f"{tma_tensor_a = !s}") + # Setup TMA load for B + b_op = utils.sm100.cluster_shape_to_tma_atom_B( + self.cluster_shape_mn, self._tiled_mma.thr_id + ) + log(f"{b_op = !s}") + # Get written of the pipeline dimension. + b_smem_layout = cute.slice_(self.b_smem_layout_staged, (None, None, None, 0)) + log(f"{b_smem_layout = !s}") + tma_atom_b, tma_tensor_b = cute.nvgpu.make_tiled_tma_atom_B( + b_op, + b, + b_smem_layout, + self.mma_tiler, + self._tiled_mma, + self.cluster_layout_vmnk.shape, + internal_type=( + cutlass.TFloat32 if b.element_type is cutlass.Float32 else None + ), + ) + log(f"{tma_atom_b = !s}") + log(f"{tma_tensor_b = !s}") + a_copy_size = cute.size_in_bytes(self.a_dtype, a_smem_layout) + b_copy_size = cute.size_in_bytes(self.b_dtype, b_smem_layout) + self.num_tma_load_bytes = (a_copy_size + b_copy_size) * atom_thr_size + log(f"{self.num_tma_load_bytes = }") + + # Set the TMA related arguments for the tensors used in the EFC. + self.efc.jit.create_tma_arguments() + + # Compute grid size + self.tile_sched_params, grid = self._compute_grid( + d, self.cta_tile_shape_mnk, self.cluster_shape_mn, max_active_clusters + ) + + self.efc.jit.create_supplemental_arguments_for_kernel() + + # Launch the kernel synchronously + self.kernel( + self._tiled_mma, + tma_atom_a, + tma_tensor_a, + tma_atom_b, + tma_tensor_b, + self.cluster_layout_vmnk, + self.a_smem_layout_staged, + self.b_smem_layout_staged, + self.epi_tile, + self.tile_sched_params, + self.efc.kernel.pack_arguments(), + ).launch( + grid=grid, + block=[self.threads_per_cta, 1, 1], + cluster=(*self.cluster_shape_mn, 1), + stream=stream, + ) + return + + # GPU device kernel + @cute.kernel + def kernel( + self, + tiled_mma: cute.TiledMma, + tma_atom_a: cute.CopyAtom, + mA_mkl: cute.Tensor, + tma_atom_b: cute.CopyAtom, + mB_nkl: cute.Tensor, + cluster_layout_vmnk: cute.Layout, + a_smem_layout_staged: cute.ComposedLayout, + b_smem_layout_staged: cute.ComposedLayout, + epi_tile: cute.Tile, + tile_sched_params: utils.PersistentTileSchedulerParams, + supplemental_parameters: Tuple, + ): + """ + GPU device kernel performing the Persistent batched GEMM computation. + """ + # Process the variadic parameters. + self.efc.kernel.unpack_parameters(supplemental_parameters) + + warp_idx = cute.arch.warp_idx() + warp_idx = cute.arch.make_warp_uniform(warp_idx) + + # + # Prefetch tma desc + # + if warp_idx == self.tma_warp_id: + cpasync.prefetch_descriptor(tma_atom_a) + cpasync.prefetch_descriptor(tma_atom_b) + # Prefetch the TMA descriptors for all the supplemental tensors. + self.efc.kernel.prefetch_tma_descriptors() + + # + # Setup cta/thread coordinates + # + # Coords inside cluster + bidx, bidy, bidz = cute.arch.block_idx() + mma_tile_coord_v = bidx % cute.size(tiled_mma.thr_id.shape) + is_leader_cta = mma_tile_coord_v == 0 + cta_rank_in_cluster = cute.arch.make_warp_uniform( + cute.arch.block_idx_in_cluster() + ) + block_in_cluster_coord_vmnk = cluster_layout_vmnk.get_flat_coord( + cta_rank_in_cluster + ) + # Coord inside cta + tidx, _, _ = cute.arch.thread_idx() + + # + # Alloc and init: a+b full/empty, accumulator full/empty, tensor memory dealloc barrier + # + self.buffer_align_bytes = 1024 + + # Define shared storage for kernel + @cute.struct + class SharedStorage: + ab_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage] + ab_empty_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage] + acc_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage] + acc_empty_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage] + # Barriers used by the supplemental load tensor pipeline. + c_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_c_stage] + c_empty_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_c_stage] + tmem_dealloc_mbar_ptr: cutlass.Int64 + tmem_holding_buf: cutlass.Int32 + # (MMA, MMA_M, MMA_K, STAGE) + sA: cute.struct.Align[ + cute.struct.MemRange[ + self.a_dtype, cute.cosize(a_smem_layout_staged.outer) + ], + self.buffer_align_bytes, + ] + # (MMA, MMA_N, MMA_K, STAGE) + sB: cute.struct.Align[ + cute.struct.MemRange[ + self.b_dtype, cute.cosize(b_smem_layout_staged.outer) + ], + self.buffer_align_bytes, + ] + + self.shared_storage = SharedStorage + + self.smem = utils.SmemAllocator() + storage = self.smem.allocate(self.shared_storage) + + # Allocate the shared memory for all the supplemental tensors. + self.efc.kernel.allocate_smem() + + # Initialize mainloop ab_pipeline (barrier) and states + ab_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) + num_tma_producer = self.num_mcast_ctas_a + self.num_mcast_ctas_b - 1 + ab_pipeline_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, num_tma_producer + ) + ab_pipeline = pipeline.PipelineTmaUmma.create( + barrier_storage=storage.ab_full_mbar_ptr.data_ptr(), + num_stages=self.num_ab_stage, + producer_group=ab_pipeline_producer_group, + consumer_group=ab_pipeline_consumer_group, + tx_count=self.num_tma_load_bytes, + cta_layout_vmnk=cluster_layout_vmnk, + ) + + # Initialize acc_pipeline (barrier) and states + acc_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) + num_acc_consumer_threads = len(self.epilogue_warp_id) * ( + 2 if self.use_2cta_instrs else 1 + ) + acc_pipeline_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, num_acc_consumer_threads + ) + acc_pipeline = pipeline.PipelineUmmaAsync.create( + barrier_storage=storage.acc_full_mbar_ptr.data_ptr(), + num_stages=self.num_acc_stage, + producer_group=acc_pipeline_producer_group, + consumer_group=acc_pipeline_consumer_group, + cta_layout_vmnk=cluster_layout_vmnk, + ) + + # Load pipeline, used to load all the supplemental tensors of the + # epilogue. + c_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) + c_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + len(self.epilogue_warp_id), + ) + c_pipeline = pipeline.PipelineTmaAsync.create( + barrier_storage=storage.c_full_mbar_ptr.data_ptr(), + num_stages=self.num_c_stage, + producer_group=c_producer_group, + consumer_group=c_consumer_group, + # Unlock the barrier when all the tensor bytes have been loaded. + tx_count=self.efc.jit.total_tma_load_bytes, + ) + + tmem_alloc_barrier = pipeline.NamedBarrier( + barrier_id=self.tmem_alloc_sync_bar_id, + num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id)), + ) + # Tensor memory dealloc barrier init + tmem = utils.TmemAllocator( + storage.tmem_holding_buf, + barrier_for_retrieve=tmem_alloc_barrier, + allocator_warp_id=self.epilogue_warp_id[0], + is_two_cta=self.use_2cta_instrs, + two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar_ptr, + ) + + # Cluster arrive after barrier init + if cute.size(self.cluster_shape_mn) > 1: + cute.arch.cluster_arrive_relaxed() + + # + # Setup smem tensor A/B + # + # (MMA, MMA_M, MMA_K, STAGE) + sA = storage.sA.get_tensor( + a_smem_layout_staged.outer, swizzle=a_smem_layout_staged.inner + ) + # (MMA, MMA_N, MMA_K, STAGE) + sB = storage.sB.get_tensor( + b_smem_layout_staged.outer, swizzle=b_smem_layout_staged.inner + ) + + # + # Compute multicast mask for A/B buffer full + # + a_full_mcast_mask = None + b_full_mcast_mask = None + if cutlass.const_expr( + self.is_a_mcast or self.is_b_mcast or self.use_2cta_instrs + ): + a_full_mcast_mask = cpasync.create_tma_multicast_mask( + cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=2 + ) + b_full_mcast_mask = cpasync.create_tma_multicast_mask( + cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=1 + ) + + # + # Local_tile partition global tensors + # + # (bM, bK, loopM, loopK, loopL) + gA_mkl = cute.local_tile( + mA_mkl, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None) + ) + # (bN, bK, loopN, loopK, loopL) + gB_nkl = cute.local_tile( + mB_nkl, cute.slice_(self.mma_tiler, (0, None, None)), (None, None, None) + ) + + k_tile_cnt = cute.size(gA_mkl, mode=[3]) + + # + # Partition global tensor for TiledMMA_A/B/D + # + self.thr_mma = tiled_mma.get_slice(mma_tile_coord_v) + # (MMA, MMA_M, MMA_K, loopM, loopK, loopL) + tCgA = self.thr_mma.partition_A(gA_mkl) + log(f"{tCgA = !s}") + # (MMA, MMA_N, MMA_K, loopN, loopK, loopL) + tCgB = self.thr_mma.partition_B(gB_nkl) + log(f"{tCgB = !s}") + # Create the local_tile gX_mnl for all the EFC supplemental tensors. + self.efc.kernel.partition_global_tensors_for_tiled_mma() + + # + # Partition global/shared tensor for TMA load A/B + # + # TMA load A partition_S/D + a_cta_layout = cute.make_layout( + cute.slice_(cluster_layout_vmnk, (0, 0, None, 0)).shape + ) + # ((atom_v, rest_v), STAGE) + # ((atom_v, rest_v), tiles_m, tiles_k, tiles_l) + tAsA, tAgA = cpasync.tma_partition( + tma_atom_a, + block_in_cluster_coord_vmnk[2], + a_cta_layout, + cute.group_modes(sA, 0, 3), + cute.group_modes(tCgA, 0, 3), + ) + # TMA load B partition_S/D + b_cta_layout = cute.make_layout( + cute.slice_(cluster_layout_vmnk, (0, None, 0, 0)).shape + ) + # ((atom_v, rest_v), STAGE) + # ((atom_v, rest_v), tiles_n, tiles_k, tiles_l) + tBsB, tBgB = cpasync.tma_partition( + tma_atom_b, + block_in_cluster_coord_vmnk[1], + b_cta_layout, + cute.group_modes(sB, 0, 3), + cute.group_modes(tCgB, 0, 3), + ) + + # + # Partition shared/tensor memory tensor for TiledMMA_A/B/C/D + # + # (MMA, MMA_M, MMA_K, STAGE) + tCrA = tiled_mma.make_fragment_A(sA) + log(f"{tCrA = !s}") + # (MMA, MMA_N, MMA_K, STAGE) + tCrB = tiled_mma.make_fragment_B(sB) + log(f"{tCrB = !s}") + # (MMA, MMA_M, MMA_N) + acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2]) + log(f"{acc_shape = !s}") + # (MMA, MMA_M, MMA_N, STAGE) + tCtAcc_fake = tiled_mma.make_fragment_C( + cute.append(acc_shape, self.num_acc_stage) + ) + log(f"{tCtAcc_fake = !s}") + + # Named barriers + # + cta_sync_barrier = pipeline.NamedBarrier( + self.cta_sync_bar_id, self.threads_per_cta + ) + epilogue_sync_barrier = pipeline.NamedBarrier( + self.epilogue_sync_bar_id, 32 * len(self.epilogue_warp_id) + ) + + # + # Cluster wait before tensor memory alloc + # + if cute.size(self.cluster_shape_mn) > 1: + cute.arch.cluster_wait() + else: + cta_sync_barrier.arrive_and_wait() + + # + # Specialized TMA load warp + # + + if warp_idx == self.tma_warp_id: + # + # Persistent tile scheduling loop + # + tile_sched = utils.StaticPersistentTileScheduler.create( + tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim() + ) + work_tile = tile_sched.initial_work_tile_info() + + ab_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.num_ab_stage + ) + + while work_tile.is_valid_tile: + # Get tile coord from tile scheduler + cur_tile_coord = work_tile.tile_idx + mma_tile_coord_mnl = ( + cur_tile_coord[0] // cute.size(tiled_mma.thr_id.shape), + cur_tile_coord[1], + cur_tile_coord[2], + ) + + # + # Slice to per mma tile index + # + # ((atom_v, rest_v), loopK) + tAgA_slice = tAgA[ + (None, mma_tile_coord_mnl[0], None, mma_tile_coord_mnl[2]) + ] + # ((atom_v, rest_v), loopK) + tBgB_slice = tBgB[ + (None, mma_tile_coord_mnl[1], None, mma_tile_coord_mnl[2]) + ] + + # Peek (try_wait) AB buffer empty for k_tile = prefetch_k_tile_cnt + ab_producer_state.reset_count() + peek_ab_empty_status = cutlass.Boolean(1) + if ab_producer_state.count < k_tile_cnt: + peek_ab_empty_status = ab_pipeline.producer_try_acquire( + ab_producer_state + ) + + # + # Tma load loop + # + for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1): + # Conditionally wait for AB buffer empty + ab_pipeline.producer_acquire( + ab_producer_state, peek_ab_empty_status + ) + + # TMA load A/B + cute.copy( + tma_atom_a, + tAgA_slice[(None, ab_producer_state.count)], + tAsA[(None, ab_producer_state.index)], + tma_bar_ptr=ab_pipeline.producer_get_barrier(ab_producer_state), + mcast_mask=a_full_mcast_mask, + ) + cute.copy( + tma_atom_b, + tBgB_slice[(None, ab_producer_state.count)], + tBsB[(None, ab_producer_state.index)], + tma_bar_ptr=ab_pipeline.producer_get_barrier(ab_producer_state), + mcast_mask=b_full_mcast_mask, + ) + + # Peek (try_wait) AB buffer empty for k_tile = prefetch_k_tile_cnt + k_tile + 1 + ab_producer_state.advance() + peek_ab_empty_status = cutlass.Boolean(1) + if ab_producer_state.count < k_tile_cnt: + peek_ab_empty_status = ab_pipeline.producer_try_acquire( + ab_producer_state + ) + + # + # Advance to next tile + # + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + + # + # Wait A/B buffer empty + # + ab_pipeline.producer_tail(ab_producer_state) + + # + # Specialized MMA warp + # + if warp_idx == self.mma_warp_id: + # + # Bar sync for retrieve tensor memory ptr from shared mem + # + tmem.wait_for_alloc() + + # + # Retrieving tensor memory ptr and make accumulator tensor + # + tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + # (MMA, MMA_M, MMA_N, STAGE) + tCtAcc_base = cute.make_tensor(tmem_ptr, tCtAcc_fake.layout) + + # + # Persistent tile scheduling loop + # + tile_sched = utils.StaticPersistentTileScheduler.create( + tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim() + ) + work_tile = tile_sched.initial_work_tile_info() + + ab_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_ab_stage + ) + acc_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.num_acc_stage + ) + + while work_tile.is_valid_tile: + # Get tile coord from tile scheduler + cur_tile_coord = work_tile.tile_idx + mma_tile_coord_mnl = ( + cur_tile_coord[0] // cute.size(tiled_mma.thr_id.shape), + cur_tile_coord[1], + cur_tile_coord[2], + ) + + # Set tensor memory buffer for current tile + # (MMA, MMA_M, MMA_N) + tCtAcc = tCtAcc_base[(None, None, None, acc_producer_state.index)] + + # Peek (try_wait) AB buffer full for k_tile = 0 + ab_consumer_state.reset_count() + peek_ab_full_status = cutlass.Boolean(1) + if ab_consumer_state.count < k_tile_cnt and is_leader_cta: + peek_ab_full_status = ab_pipeline.consumer_try_wait( + ab_consumer_state + ) + + # + # Wait for accumulator buffer empty + # + if is_leader_cta: + acc_pipeline.producer_acquire(acc_producer_state) + + # + # Reset the ACCUMULATE field for each tile + # + tiled_mma.set(tcgen05.Field.ACCUMULATE, False) + + # + # Mma mainloop + # + for k_tile in range(k_tile_cnt): + if is_leader_cta: + # Conditionally wait for AB buffer full + ab_pipeline.consumer_wait( + ab_consumer_state, peek_ab_full_status + ) + + # tCtAcc += tCrA * tCrB + num_k_blocks = cute.size(tCrA, mode=[2]) + for k_block_idx in cutlass.range( + num_k_blocks, unroll_full=True + ): + k_block_coord = ( + None, + None, + k_block_idx, + ab_consumer_state.index, + ) + + cute.gemm( + tiled_mma, + tCtAcc, + tCrA[k_block_coord], + tCrB[k_block_coord], + tCtAcc, + ) + # Enable accumulate on tCtAcc after first k_block + tiled_mma.set(tcgen05.Field.ACCUMULATE, True) + + # Async arrive AB buffer empty + ab_pipeline.consumer_release(ab_consumer_state) + + # Peek (try_wait) AB buffer full for k_tile = k_tile + 1 + ab_consumer_state.advance() + peek_ab_full_status = cutlass.Boolean(1) + if ab_consumer_state.count < k_tile_cnt: + if is_leader_cta: + peek_ab_full_status = ab_pipeline.consumer_try_wait( + ab_consumer_state + ) + + # + # Async arrive accumulator buffer full + # + if is_leader_cta: + acc_pipeline.producer_commit(acc_producer_state) + acc_producer_state.advance() + + # + # Advance to next tile + # + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + + # + # Wait for accumulator buffer empty + # + acc_pipeline.producer_tail(acc_producer_state) + # + # Specialized epilogue warps + # + if warp_idx < self.mma_warp_id: + # + # Alloc tensor memory buffer + # + tmem.allocate(self.num_tmem_alloc_cols) + + # + # Bar sync for retrieve tensor memory ptr from shared memory + # + tmem.wait_for_alloc() + + # + # Retrieving tensor memory ptr and make accumulator tensor + # + tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + log(f"tmem_ptr = {tmem_ptr!s}") + # (MMA, MMA_M, MMA_N, STAGE) + tCtAcc_base = cute.make_tensor(tmem_ptr, tCtAcc_fake.layout) + log(f"tCtAcc_base = {tCtAcc_base!s}") + # + # Partition for epilogue + # + epi_tidx = tidx + tCgD = self.efc.kernel.tCgD_written[self.d_name_bigger] + log(f"tCgD (aka tCgD_written[{self.d_name_bigger}])= {tCgD!s}") + + ( + tiled_copy_t2r, # (EPI_TILE_M, EPI_TILE_N) + tTR_tAcc_base, # (T2R, T2R_M, T2R_N, EPI_M, EPI_M, STAGE) + tTR_rAcc, # (T2R, T2R_M, T2R_N) + ) = self.epilogue_tmem_copy_and_partition( + epi_tidx, tCtAcc_base, tCgD, epi_tile + ) + log(f"{tiled_copy_t2r = !s}") + log(f"{tTR_tAcc_base = !s}") + log(f"{tTR_rAcc = !s}") + # Copy and partition for the supplemental EFC tensors. + self.efc.kernel.copy_and_partition_supplemental_rmem_tensors( + tiled_copy_t2r, tTR_rAcc, epi_tidx, epi_tile + ) + + # + # Persistent tile scheduling loop + # + tile_sched = utils.StaticPersistentTileScheduler.create( + tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim() + ) + work_tile = tile_sched.initial_work_tile_info() + + acc_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_acc_stage + ) + + # Store D pipeline used for all the written tensors in the epilogue. + d_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + 32 * len(self.epilogue_warp_id), + ) + d_pipeline = pipeline.PipelineTmaStore.create( + num_stages=self.num_d_stage, + producer_group=d_producer_group, + ) + + c_pipeline_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_c_stage + ) + + while work_tile.is_valid_tile: + # Get tile coord from tile scheduler + cur_tile_coord = work_tile.tile_idx + mma_tile_coord_mnl = ( + cur_tile_coord[0] // cute.size(tiled_mma.thr_id.shape), + cur_tile_coord[1], + cur_tile_coord[2], + ) + + # Slice the supplemental written tensors per MMA tile index. + self.efc.kernel.slice_written_tensors_per_mma_tile_index( + mma_tile_coord_mnl + ) + + # Set tensor memory buffer for current tile + # (T2R, T2R_M, T2R_N, EPI_M, EPI_M) + tTR_tAcc = tTR_tAcc_base[ + (None, None, None, None, None, acc_consumer_state.index) + ] + log(f"tTR_tAcc = {tTR_tAcc!s}") + # + # Wait for accumulator buffer full + # + acc_pipeline.consumer_wait(acc_consumer_state) + + # Group together the EPI_M, EPI_M which are starting at group 3. + # (T2R, T2R_M, T2R_N, (EPI_M, EPI_M)) + tTR_tAcc = cute.group_modes(tTR_tAcc, 3, cute.rank(tTR_tAcc)) + log(f"group_modes tTR_tAcc = {tTR_tAcc!s}") + # + # Store accumulator to global memory in subtiles + # + # Use EPI_M*EPI_M to iterate using the 1-D coordinate. + subtile_cnt = cute.size(tTR_tAcc.shape, mode=[3]) + num_prev_subtiles = tile_sched.num_tiles_executed * subtile_cnt + for subtile_idx in cutlass.range(subtile_cnt): + # + # Load accumulator from tensor memory buffer to register + # + tTR_tAcc_mn = tTR_tAcc[(None, None, None, subtile_idx)] + cute.copy(tiled_copy_t2r, tTR_tAcc_mn, tTR_rAcc) + log(f"cute.copy tiled_copy_t2r = {tiled_copy_t2r!s}") + log(f"cute.copy tTR_tAcc_mn = {tTR_tAcc_mn!s}") + log(f"cute.copy tTR_rAcc = {tTR_rAcc!s}") + + # Wait for the EFC tensor loads to complete. + if cutlass.const_expr(self.efc.read_tensor_names): + # The wait is blocking even if the tx_count is 0 when + # there is no tensor to load. So need to skip it when + # there is no tensor to read. + c_pipeline.consumer_wait(c_pipeline_consumer_state) + + # Load supplemental tensors from shared memory to register. + self.efc.kernel.load_tensors_from_smem_to_register( + c_pipeline_consumer_state.index + ) + + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) + c_pipeline.consumer_release(c_pipeline_consumer_state) + + # Advance pipeline states + c_pipeline_consumer_state.advance() + + # + # Perform epilogue op on accumulator. + # + tiled_copy_r2s = self.efc.kernel.tiled_copy_r2s[self.d_name_bigger] + log(f"tiled_copy_r2s = {tiled_copy_r2s!s}") + # Use a SimpleNamespace to pass easily some local content as + # an extensible class compatible with CuTe DSL + # implementation. + epilogue_context = types.SimpleNamespace() + # Load the accumulator cast to the epi_dtype used to do all + # the computations in the epilogue. + # Retile the accumulator subtile to fit the destination + # subtile vector TV layout. + epilogue_context.acc_vec = ( + tiled_copy_r2s.retile(tTR_rAcc).load().to(self.epi_dtype) + ) + log(f"before .retile tTR_rAcc = {tTR_rAcc!s}") + log( + f"tiled_copy_r2s.retile(tTR_rAcc) = {tiled_copy_r2s.retile(tTR_rAcc)!s}" + ) + log( + f"tiled_copy_r2s.retile(tTR_rAcc).load() = {tiled_copy_r2s.retile(tTR_rAcc).load()!s}" + ) + log(f"epilogue_context.acc_vec = {epilogue_context.acc_vec!s}") + + # Execute the EFC epilogue. + self.efc.kernel.epilogue_computation(epilogue_context) + d_buffer = (num_prev_subtiles + subtile_idx) % self.num_d_stage + + # Store the EFC written tensors to shared memory. + self.efc.kernel.store_written_tensors_to_smem(d_buffer) + + # Fence and barrier to make sure shared memory store is visible to TMA store + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) + epilogue_sync_barrier.arrive_and_wait() + + # + # TMA store D to global memory + # + if warp_idx == self.epilogue_warp_id[0]: + # Store with TMA the written EFC tensors to global memory. + self.efc.kernel.tma_store_written_tensors_to_gmem( + d_buffer, subtile_idx + ) + # Fence and barrier to make sure shared memory store is visible to TMA store + d_pipeline.producer_commit() + d_pipeline.producer_acquire() + + epilogue_sync_barrier.arrive_and_wait() + + # + # Async arrive accumulator buffer empty + # + with cute.arch.elect_one(): + acc_pipeline.consumer_release(acc_consumer_state) + acc_consumer_state.advance() + + # + # Advance to next tile + # + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + + # + # Dealloc the tensor memory buffer + # + tmem.relinquish_alloc_permit() + epilogue_sync_barrier.arrive_and_wait() + tmem.free(tmem_ptr) + # + # Wait for D store complete + # + d_pipeline.producer_tail() + + # + # Specialized epilog load warp + # + if warp_idx == self.epilogue_load_warp_id: + # Create the tiled tensors to be loaded in the epilogue. + self.efc.kernel.create_epilogue_subtile_tensors(tidx, epi_tile) + + tile_sched = utils.StaticPersistentTileScheduler.create( + tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim() + ) + work_tile = tile_sched.initial_work_tile_info() + + c_pipeline_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.num_c_stage + ) + + # Setup the pipelines reading the EFC supplemental tensors. + + while work_tile.is_valid_tile: + # Get tile coord from tile scheduler + cur_tile_coord = work_tile.tile_idx + mma_tile_coord_mnl = ( + cur_tile_coord[0] // cute.size(tiled_mma.thr_id.shape), + cur_tile_coord[1], + cur_tile_coord[2], + ) + # Prepare the EFC tensors to be loaded by the subtiles. + subtile_cnt = self.efc.kernel.prepare_tensor_load_for_subtiles( + mma_tile_coord_mnl, + ) + + # Assume the pipeline can work even in the case there is no + # tensor to load and so subtile_cnt is 0. + for subtile_idx in cutlass.range(subtile_cnt): + # Load C from global memory to shared memory. + c_pipeline.producer_acquire(c_pipeline_producer_state) + + # Load the subtiles of EFC tensors. + self.efc.kernel.load_tensor_subtiles( + subtile_idx, c_pipeline, c_pipeline_producer_state + ) + + c_pipeline_producer_state.advance() + + # + # Advance to next tile + # + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + + # + # Wait for the load buffer to be empty. + # + c_pipeline.producer_tail(c_pipeline_producer_state) + + def epilogue_tmem_copy_and_partition( + self, + tidx: cutlass.Int32, + tAcc: cute.Tensor, + tCgC: cute.Tensor, + epi_tile: cute.Tile, + ) -> Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]: + """ + Make tiledCopy for tensor memory load, then use it to partition tensor memory (source) and register array (destination). + """ + # Make tiledCopy for tensor memory load + copy_atom_t2r = sm100_utils.get_tmem_load_op( + self.cta_tile_shape_mnk, + self.d_layout, # Take this as the reference layout for the epilogue tile. + self.epi_dtype, # But we get the accumulator as epi_dtype in the epilogue. + self.acc_dtype, + epi_tile, + self.use_2cta_instrs, + ) + + # (EPI_TILE_M, EPI_TILE_N, EPI_M, EPI_N, STAGE) + tAcc_epi = cute.flat_divide( + tAcc[((None, None), 0, 0, None)], + epi_tile, + ) + # (EPI_TILE_M, EPI_TILE_N) + tiled_copy_t2r = tcgen05.make_tmem_copy( + copy_atom_t2r, tAcc_epi[(None, None, 0, 0, 0)] + ) + thr_copy_t2r = tiled_copy_t2r.get_slice(tidx) + # (T2R, T2R_M, T2R_N, EPI_M, EPI_M, STAGE) + tTR_tAcc = thr_copy_t2r.partition_S(tAcc_epi) + + # (EPI_TILE_M, EPI_TILE_N, EPI_M, EPI_N, RestM, RestN, RestL) + tCgC_epi = cute.flat_divide( + tCgC[((None, None), 0, 0, None, None, None)], epi_tile + ) + # (T2R, T2R_M, T2R_N, EPI_M, EPI_N, RestM, RestN, RestL) + tTR_gC = thr_copy_t2r.partition_D(tCgC_epi) + # (T2R, T2R_M, T2R_N) + tTR_rAcc = cute.make_rmem_tensor( + tTR_gC[(None, None, None, 0, 0, 0, 0, 0)].shape, self.acc_dtype + ) + return tiled_copy_t2r, tTR_tAcc, tTR_rAcc + + def epilogue_smem_copy_and_partition_load( + self, + tiled_copy_t2r: cute.TiledCopy, + tTR_rC: cute.Tensor, + tidx: cutlass.Int32, + sC: cute.Tensor, + ) -> Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]: + """ + Make tiledCopy for shared memory load, then use it to partition register array (destination) and shared memory (source). + + :param tiled_copy_t2r: The tiled copy operation for tmem to register copy(t2r) + :type tiled_copy_t2r: cute.TiledCopy + :param tTR_rC: The partitioned accumulator tensor + :type tTR_rC: cute.Tensor + :param tidx: The thread index in epilogue warp groups + :type tidx: cutlass.Int32 + :param sC: The shared memory tensor to be copied and partitioned + :type sC: cute.Tensor + + :return: A tuple containing (tiled_copy_s2r, tSR_rC, tSR_sC) where: + - tiled_copy_s2r: The tiled copy operation for smem to register copy(s2r) + - tSR_rC: The partitioned tensor C (register destination) + - tSR_sC: The partitioned tensor C (smem source) + :rtype: Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor] + """ + copy_atom_s2r = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), self.c_dtype) + tiled_copy_s2r = cute.make_tiled_copy_D(copy_atom_s2r, tiled_copy_t2r) + # (S2R, S2R_M, S2R_N, PIPE_C) + thr_copy_s2r = tiled_copy_s2r.get_slice(tidx) + tSR_sC = thr_copy_s2r.partition_D(sC) + # (S2R, S2R_M, S2R_N) + tSR_rC = tiled_copy_s2r.retile(tTR_rC) + return tiled_copy_s2r, tSR_rC, tSR_sC + + def epilogue_gmem_copy_and_partition( + self, + tidx: cutlass.Int32, + atom: Union[cute.CopyAtom, cute.TiledCopy], + gC_mnl: cute.Tensor, + epi_tile: cute.Tile, + sC: cute.Tensor, + dtype: Type[cutlass.Numeric], + ) -> Tuple[cute.CopyAtom, cute.Tensor, cute.Tensor]: + """Make tiledCopy for global memory store, then use it to: + - partition register array (source) and global memory (destination) for none TMA store version; + - partition shared memory (source) and global memory (destination) for TMA store version. + + :param tidx: The thread index in epilogue warp groups + :type tidx: cutlass.Int32 + :param atom: The copy_atom_c to be used for TMA store version, or tiled_copy_t2r for none TMA store version + :type atom: cute.CopyAtom or cute.TiledCopy + :param gC_mnl: The global tensor C + :type gC_mnl: cute.Tensor + :param epi_tile: The epilogue tiler + :type epi_tile: cute.Tile + :param sC: The shared memory tensor to be copied and partitioned + :type sC: cute.Tensor + + :return: A tuple containing either: + - For TMA store: (tma_atom_c, bSG_sC, bSG_gC) where: + - tma_atom_c: The TMA copy atom + - bSG_sC: The partitioned shared memory tensor C + - bSG_gC: The partitioned global tensor C + - For non-TMA store: (simt_atom, tTR_rC, tTR_gC) where: + - simt_atom: The SIMT copy atom + - tTR_rC: The register tensor C + - tTR_gC: The partitioned global tensor C + :rtype: Tuple[cute.CopyAtom, cute.Tensor, cute.Tensor] + """ + # (EPI_TILE_M, EPI_TILE_N, EPI_M, EPI_N, tiles_m, tiles_n, tiles_l) + gC_epi = cute.flat_divide( + gC_mnl[((None, None), 0, 0, None, None, None)], epi_tile + ) + + tma_atom_c = atom + sC_for_tma_partition = cute.group_modes(sC, 0, 2) + gC_for_tma_partition = cute.group_modes(gC_epi, 0, 2) + # ((ATOM_V, REST_V), EPI_M, EPI_N) + # ((ATOM_V, REST_V), EPI_M, EPI_N, tiles_m, tiles_n, tiles_l) + bSG_sC, bSG_gC = cpasync.tma_partition( + tma_atom_c, + 0, + cute.make_layout(1), + sC_for_tma_partition, + gC_for_tma_partition, + ) + return tma_atom_c, bSG_sC, bSG_gC + + def compute_stages(self) -> None: + """Compute and set the number of stages for A/B/C/D operands. + + Uses instance attributes to compute and assign: + `self.num_acc_stage`, `self.num_ab_stage`, `self.num_c_stage`, + and `self.num_d_stage`. + """ + # Defaults + self.num_acc_stage = 2 + # To read the tensors needed for the epilogue: + self.num_c_stage = 2 + # To write the tensors produced by the epilogue: + self.num_d_stage = 2 + + # Calculate smem layout and size for one stage of A, B, C, and D + a_smem_layout_stage_one = sm100_utils.make_smem_layout_a( + self._tiled_mma, self.mma_tiler, self.a_dtype, 1 + ) + b_smem_layout_staged_one = sm100_utils.make_smem_layout_b( + self._tiled_mma, self.mma_tiler, self.b_dtype, 1 + ) + + # Get the contribution from the tensors used in the EFC. + self.efc.jit.compute_stage() + + ab_bytes_per_stage = cute.size_in_bytes( + self.a_dtype, a_smem_layout_stage_one + ) + cute.size_in_bytes(self.b_dtype, b_smem_layout_staged_one) + mbar_helpers_bytes = 1024 + # Contribution from the tensors loaded in the EFC. + c_bytes_per_stage = self.efc.jit.smem_size_in_bytes_of_read_tensors() + c_bytes = c_bytes_per_stage * self.num_c_stage + # Contribution from the tensors stored in the EFC. There is at least 1 + # written tensor, so the following is strictly positive. + d_bytes_per_stage = self.efc.jit.smem_size_in_bytes_of_written_tensors() + d_bytes = d_bytes_per_stage * self.num_d_stage + + # Calculate A/B stages + self.num_ab_stage = ( + self.smem_capacity // self.occupancy + - (mbar_helpers_bytes + c_bytes + d_bytes) + ) // ab_bytes_per_stage + log(f"\t{self.num_ab_stage = }") + + if self.num_ab_stage <= 0: + raise MemoryError("Not enough smem capacity to allocate all the tensors.") + + # Refine epilogue stages: + # Calculate remaining smem after allocating for A/B stages and reserved bytes. + # Add remaining unused smem to epilogue. + self.num_d_stage += ( + self.smem_capacity + - self.occupancy * ab_bytes_per_stage * self.num_ab_stage + - self.occupancy * (mbar_helpers_bytes + c_bytes + d_bytes) + ) // (self.occupancy * d_bytes_per_stage) + log(f"\tnew {self.num_d_stage = }") + + @staticmethod + def _compute_grid( + d: cute.Tensor, + cta_tile_shape_mnk: Tuple[int, int, int], + cluster_shape_mn: Tuple[int, int], + max_active_clusters: cutlass.Constexpr, + ) -> Tuple[utils.PersistentTileSchedulerParams, Tuple[int, int, int]]: + """Use persistent tile scheduler to compute the grid size for the output tensor D. + + :param d: The output tensor D + :type d: cute.Tensor + :param cta_tile_shape_mnk: The shape (M, N, K) of the CTA tile. + :type cta_tile_shape_mnk: tuple[int, int, int] + :param cluster_shape_mn: Shape of each cluster in M, N dimensions. + :type cluster_shape_mn: tuple[int, int] + :param max_active_clusters: Maximum number of active clusters. + :type max_active_clusters: cutlass.Constexpr + + :return: A tuple containing: + - tile_sched_params: Parameters for the persistent tile scheduler. + - grid: Grid shape for kernel launch. + :rtype: Tuple[utils.PersistentTileSchedulerParams, tuple[int, int, int]] + """ + log(f"compute_grid: {max_active_clusters = }") + d_shape = cute.slice_(cta_tile_shape_mnk, (None, None, 0)) + gd = cute.zipped_divide(d, tiler=d_shape) + num_ctas_mnl = gd[(0, (None, None, None))].shape + common_efc.if_debug( + lambda: cute.printf("compute_grid: num_ctas_mnl = {}", num_ctas_mnl) + ) + cluster_shape_mnl = (*cluster_shape_mn, 1) + + tile_sched_params = utils.PersistentTileSchedulerParams( + num_ctas_mnl, cluster_shape_mnl + ) + grid = utils.StaticPersistentTileScheduler.get_grid_shape( + tile_sched_params, max_active_clusters + ) + common_efc.if_debug(lambda: cute.printf("compute_grid: grid = {}", grid)) + return tile_sched_params, grid + + def compute_num_tmem_alloc_cols(self) -> None: + """Compute and set the number of tensor memory allocation columns. + + This method uses the instance attributes computed during setup to + determine the number of tensor memory allocation columns and stores + the result in `self.num_tmem_alloc_cols`. + """ + acc_shape = self._tiled_mma.partition_shape_C(self.mma_tiler[:2]) + log(f"compute_num_tmem_alloc_cols: {acc_shape = !s}") + tCtAcc_fake = self._tiled_mma.make_fragment_C( + cute.append(acc_shape, self.num_acc_stage) + ) + log(f"compute_num_tmem_alloc_cols: {tCtAcc_fake = !s}") + self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols(tCtAcc_fake) + log(f"compute_num_tmem_alloc_cols: {self.num_tmem_alloc_cols = }") + + def check_valid_dtypes( + self, + ab_dtype: Type[cutlass.Numeric], + ): + """ + Check if the dtypes are valid + + :param ab_dtype: The data type of the A and B operands + :type ab_dtype: Type[cutlass.Numeric] + :param acc_dtype: The data type of the accumulator + :type acc_dtype: Type[cutlass.Numeric] + + :raises ValueError: If the dtypes are invalid or incompatible + """ + valid_ab_dtypes = { + cutlass.Float16, + cutlass.BFloat16, + cutlass.TFloat32, + cutlass.Uint8, + cutlass.Int8, + cutlass.Float8E4M3FN, + cutlass.Float8E5M2, + } + if ab_dtype not in valid_ab_dtypes: + raise ValueError( + f"Invalid A/B dtype: {ab_dtype}. " + f"Supported dtypes: {', '.join(str(dt) for dt in valid_ab_dtypes)}" + ) + + valid_acc_dtypes = {cutlass.Float32, cutlass.Float16, cutlass.Int32} + if self.acc_dtype not in valid_acc_dtypes: + raise ValueError( + f"Invalid accumulator dtype: {self.acc_dtype}. " + f"Supported dtypes: {', '.join(str(dt) for dt in valid_acc_dtypes)}" + ) + + # Define compatibility mapping between accumulator type and AB type + acc_ab_compatibility = { + cutlass.Float32: { + cutlass.Float16, + cutlass.BFloat16, + cutlass.TFloat32, + cutlass.Float8E4M3FN, + cutlass.Float8E5M2, + }, # Float32 accumulator supports floating point AB types only + cutlass.Float16: { + cutlass.Float16, + cutlass.Float8E4M3FN, + cutlass.Float8E5M2, + }, + cutlass.Int32: {cutlass.Uint8, cutlass.Int8}, + } + # Check compatibility between accumulator type and AB type + if ab_dtype not in acc_ab_compatibility[self.acc_dtype]: + compatible_types = acc_ab_compatibility[self.acc_dtype] + raise ValueError( + f"Incompatible dtype combination: A/B dtype {ab_dtype} is not compatible " + f"with accumulator dtype {self.acc_dtype}. " + f"Compatible A/B dtypes for {self.acc_dtype}: {', '.join(str(dt) for dt in compatible_types)}" + ) + + def check_valid_mma_tiler_and_cluster_shape(self): + """Check if the mma tiler and cluster shape are valid. + + :raises ValueError: If the mma tiler or cluster shape is invalid + """ + # Check invalid mma tile shape M dimension + if not ( + (not self.use_2cta_instrs and self.mma_tiler_mn[0] in [64, 128]) + or (self.use_2cta_instrs and self.mma_tiler_mn[0] in [128, 256]) + ): + expected = [128, 256] if self.use_2cta_instrs else [64, 128] + raise ValueError( + f"Invalid MMA tile M dimension: {self.mma_tiler_mn[0]}. " + f"Expected one of {expected} (use_2cta_instrs={self.use_2cta_instrs})" + ) + + # Check invalid mma tile shape N dimension + if self.mma_tiler_mn[1] not in range(32, 257, 32): + raise ValueError( + f"Invalid MMA tile N dimension: {self.mma_tiler_mn[1]}. " + f"Expected a multiple of 32 in range [32, 256]" + ) + # Check illegal cluster shape M dimension + if self.cluster_shape_mn[0] % (2 if self.use_2cta_instrs else 1) != 0: + divisor = 2 if self.use_2cta_instrs else 1 + raise ValueError( + f"Invalid cluster shape M dimension: {self.cluster_shape_mn[0]}. " + f"Must be divisible by {divisor} (use_2cta_instrs={self.use_2cta_instrs})" + ) + + def is_power_of_2(x): + return x > 0 and (x & (x - 1)) == 0 + + # Check invalid cluster shape constraints + if self.cluster_shape_mn[0] * self.cluster_shape_mn[1] > 16: + raise ValueError( + f"Invalid cluster shape: {self.cluster_shape_mn}. " + f"Product {self.cluster_shape_mn[0]} * {self.cluster_shape_mn[1]} = " + f"{self.cluster_shape_mn[0] * self.cluster_shape_mn[1]} exceeds maximum of 16" + ) + if self.cluster_shape_mn[0] <= 0 or self.cluster_shape_mn[1] <= 0: + raise ValueError( + f"Invalid cluster shape: {self.cluster_shape_mn}. " + f"Both dimensions must be positive" + ) + if not is_power_of_2(self.cluster_shape_mn[0]) or not is_power_of_2( + self.cluster_shape_mn[1] + ): + raise ValueError( + f"Invalid cluster shape: {self.cluster_shape_mn}. " + f"Both dimensions must be powers of 2" + ) + + def check_valid_tensor_alignment( + self, + m: int, + n: int, + k: int, + l: int, + ab_dtype: Type[cutlass.Numeric], + d_dtype: Type[cutlass.Numeric], + a_major: str, + b_major: str, + cd_major: str, + ): + """ + Check if the tensor alignment is valid + + :param m: The number of rows in the A tensor + :type m: int + :param n: The number of columns in the B tensor + :type n: int + :param k: The number of columns in the A tensor + :type k: int + :param l: The number of columns in the C tensor + :type l: int + :param ab_dtype: The data type of the A and B operands + :type ab_dtype: Type[cutlass.Numeric] + :param d_dtype: The data type of the D tensor + :type d_dtype: Type[cutlass.Numeric] + :param a_major: The major axis of the A tensor + :type a_major: str + :param b_major: The major axis of the B tensor + :type b_major: str + :param cd_major: The major axis of the C/D tensor + :type cd_major: str + + :raises ValueError: If the tensor alignment is invalid + """ + + def check_contigous_16B_alignment( + dtype, is_mode0_major, tensor_shape, tensor_name + ): + major_mode_idx = 0 if is_mode0_major else 1 + num_major_elements = tensor_shape[major_mode_idx] + num_contiguous_elements = 16 * 8 // dtype.width + if num_major_elements % num_contiguous_elements != 0: + raise ValueError( + f"Invalid alignment for tensor {tensor_name}. " + f"Major dimension has {num_major_elements} elements, " + f"but requires alignment to {num_contiguous_elements} elements (16 bytes). " + f"Dtype: {dtype}, width: {dtype.width} bits" + ) + + check_contigous_16B_alignment(ab_dtype, a_major == "m", (m, k, l), "A") + check_contigous_16B_alignment(ab_dtype, b_major == "n", (n, k, l), "B") + check_contigous_16B_alignment(d_dtype, cd_major == "m", (m, n, l), "D") + + def check_implementable(self, a: cute.Tensor, b: cute.Tensor, d: cute.Tensor): + """Check if the given tensors can be implemented by this kernel. + + :param a: Input tensor A + :type a: cute.Tensor + :param b: Input tensor B + :type b: cute.Tensor + :param d: One of the tensor used as some output + :type d: cute.Tensor + + :raises CantImplementError: If the configuration is not implementable + """ + m, n, k, l = a.shape[0], b.shape[0], a.shape[1], a.shape[2] + + # infer a_major, b_major, cd_major + is_m_major_a = utils.LayoutEnum.from_tensor(a).is_m_major_a() + is_n_major_b = utils.LayoutEnum.from_tensor(b).is_n_major_b() + is_m_major_d = utils.LayoutEnum.from_tensor(d).is_m_major_c() + a_major = "m" if is_m_major_a else "k" + b_major = "n" if is_n_major_b else "k" + cd_major = "m" if is_m_major_d else "n" + + try: + # Check dtypes (raises ValueError if invalid) + self.check_valid_dtypes(a.element_type) + + # Check mma tile shape and cluster shape (raises ValueError if invalid) + self.check_valid_mma_tiler_and_cluster_shape() + + # Check problem shape for load/store alignment (raises ValueError if invalid) + self.check_valid_tensor_alignment( + m, + n, + k, + l, + a.element_type, + d.element_type, + a_major, + b_major, + cd_major, + ) + except ValueError as e: + raise cute.testing.CantImplementError(f"Configuration error: {e}") + + class CLIParser: + """Parse command-line arguments for the Blackwell Dense GEMM example.""" + + def __init__(self): + self.parser = argparse.ArgumentParser( + description="Example of Dense Persistent GEMM on Blackwell." + ) + self.parser.add_argument( + "--mnkl", + type=self.parse_comma_separated_ints, + default=(256, 256, 512, 1), + help="mnkl dimensions (comma-separated)", + ) + self.parser.add_argument( + "--mma_tiler_mn", + type=self.parse_comma_separated_ints, + default=(128, 128), + help="Mma tile shape (comma-separated)", + ) + self.parser.add_argument( + "--cluster_shape_mn", + type=self.parse_comma_separated_ints, + default=(1, 1), + help="Cluster shape (comma-separated)", + ) + self.parser.add_argument( + "--ab_dtype", type=cutlass.dtype, default=cutlass.TFloat32 + ) + self.parser.add_argument( + "--acc_dtype", type=cutlass.dtype, default=cutlass.Float32 + ) + self.parser.add_argument( + "--epi_dtype", type=cutlass.dtype, default=cutlass.Float32 + ) + self.parser.add_argument( + "--use_2cta_instrs", + action="store_true", + help="Enable 2CTA MMA instructions feature", + ) + self.parser.add_argument( + "--a_major", choices=["k", "m"], type=str, default="k" + ) + self.parser.add_argument( + "--b_major", choices=["k", "n"], type=str, default="k" + ) + self.parser.add_argument( + "--cd_major", choices=["n", "m"], type=str, default="n" + ) + self.parser.add_argument( + "--tolerance", + type=float, + default=1e-01, + help="Tolerance for validation", + ) + self.parser.add_argument( + "--warmup_iterations", type=int, default=0, help="Warmup iterations" + ) + self.parser.add_argument( + "--iterations", + type=int, + default=1, + help="Number of iterations to run the kernel", + ) + self.parser.add_argument( + "--skip_ref_check", action="store_true", help="Skip reference checking" + ) + + # A children class may add more things to parse. + self.more_parsing() + + def parse(self): + """Parse the command-line arguments.""" + args = self.parser.parse_args() + + if len(args.mnkl) != 4: + self.parser.error("--mnkl must contain exactly 4 values") + + if len(args.mma_tiler_mn) != 2: + self.parser.error("--mma_tiler_mn must contain exactly 2 values") + + if len(args.cluster_shape_mn) != 2: + self.parser.error("--cluster_shape_mn must contain exactly 2 values") + + return args + + @staticmethod + def parse_comma_separated_ints(s: str) -> Tuple[int, ...]: + try: + return tuple(int(x.strip()) for x in s.split(",")) + except ValueError: + raise argparse.ArgumentTypeError( + "Invalid format. Expected comma-separated integers." + ) + + def more_parsing(self): + """To override to add more stuff to self.parser""" + + @staticmethod + def dtype_name(dtype: Type[cutlass.Numeric]) -> str: + """Convert a CUTLASS dtype object to its clean string name. + + This is needed to format dtype objects into CLI arguments without + full module paths. CUTLASS dtype objects have different representations: + some have a __name__ attribute while others need string parsing. + + We want "Float16" not "cutlass.Float16" or "". + + :param dtype: CUTLASS numeric data type + :return: Clean type name string (e.g., "Float16", "BFloat16") + + Example: + >>> DenseGemmEFC.dtype_name(cutlass.Float16) + 'Float16' + """ + return ( + dtype.__name__ if hasattr(dtype, "__name__") else str(dtype).split(".")[-1] + ) + + @staticmethod + def format_as_cli_args( + script_name: str, + mnkl: Tuple[int, int, int, int], + ab_dtype: Type[cutlass.Numeric], + acc_dtype: Type[cutlass.Numeric], + epi_dtype: Type[cutlass.Numeric], + a_major: str, + b_major: str, + cd_major: str, + mma_tiler_mn: Tuple[int, int], + cluster_shape_mn: Tuple[int, int], + use_2cta_instrs: bool, + tolerance: float, + ) -> str: + """Format common test parameters as CLI arguments base. + + This method formats the common parameters shared across different GEMM examples. + Subclass-specific parameters (like alpha, beta, etc.) should be added by overriding methods. + + :param script_name: Name of the Python script + :param mnkl: Matrix dimensions (M, N, K, L) + :param ab_dtype: Data type for A and B matrices + :param acc_dtype: Data type for accumulation + :param epi_dtype: Data type for epilogue + :param a_major: Major order for matrix A + :param b_major: Major order for matrix B + :param cd_major: Major order for matrices C and D + :param mma_tiler_mn: MMA tiler dimensions (M, N) + :param cluster_shape_mn: Cluster shape (M, N) + :param use_2cta_instrs: Whether to use 2CTA instructions + :param tolerance: Tolerance for validation + :return: Formatted CLI command string + """ + # Format tuples as comma-separated values + mnkl_str = ",".join(map(str, mnkl)) + mma_tiler_str = ",".join(map(str, mma_tiler_mn)) + cluster_shape_str = ",".join(map(str, cluster_shape_mn)) + + cmd = ( + f"python {script_name} " + f"--mnkl {mnkl_str} " + f"--ab_dtype {DenseGemmEFC.dtype_name(ab_dtype)} " + f"--acc_dtype {DenseGemmEFC.dtype_name(acc_dtype)} " + f"--epi_dtype {DenseGemmEFC.dtype_name(epi_dtype)} " + f"--a_major {a_major} " + f"--b_major {b_major} " + f"--cd_major {cd_major} " + f"--mma_tiler_mn {mma_tiler_str} " + f"--cluster_shape_mn {cluster_shape_str} " + f"{'--use_2cta_instrs ' if use_2cta_instrs else ''}" + f"--tolerance {tolerance}" + ) + return cmd + + def create_arguments(self, l, m, n, k, a_major, b_major, cd_major, ab_dtype): + """Create base GEMM input tensors A and B. + + This method creates the input matrices for GEMM computation. Subclasses + typically override this method to create additional supplemental tensors + for the epilogue. + + :param l: Batch dimension + :param m: M dimension (rows of A and output) + :param n: N dimension (columns of B and output) + :param k: K dimension (inner dimension) + :param a_major: Major order for A matrix ('m' or 'k') + :param b_major: Major order for B matrix ('n' or 'k') + :param cd_major: Major order for supplemental tensors ('m' or 'n') + :param ab_dtype: Data type for A and B matrices + :return: Tuple of (a_tensor, b_tensor, a_torch_cpu, b_torch_cpu) + """ + torch.manual_seed(1111) + + a_torch_cpu = cutlass_torch.matrix(l, m, k, a_major == "m", ab_dtype) + b_torch_cpu = cutlass_torch.matrix(l, n, k, b_major == "n", ab_dtype) + + a_tensor, _ = cutlass_torch.cute_tensor_like( + a_torch_cpu, ab_dtype, is_dynamic_layout=True, assumed_align=16 + ) + b_tensor, _ = cutlass_torch.cute_tensor_like( + b_torch_cpu, ab_dtype, is_dynamic_layout=True, assumed_align=16 + ) + + return ( + a_tensor, + b_tensor, + a_torch_cpu, + b_torch_cpu, + ) + + def evaluate_on_cpu( + self, + a_torch_cpu, + b_torch_cpu, + epi_dtype, + *epilogue_args, + ): + """Evaluate the GEMM and epilogue computation on CPU for validation. + + Computes the reference result by performing A*B using einsum, then + evaluates the epilogue function with the accumulator and supplemental + arguments. This updates any output tensors in epilogue_args. + + :param a_torch_cpu: Input matrix A on CPU + :param b_torch_cpu: Input matrix B on CPU + :param epi_dtype: Data type for epilogue computation + :param epilogue_args: Supplemental arguments for the epilogue (tensors and scalars) + """ + # Compute reference result + ref = torch.einsum( + "mkl,nkl->mnl", + a_torch_cpu.to(dtype=torch.float32), + b_torch_cpu.to(dtype=torch.float32), + ) + + self.efc.evaluate_on_cpu(ref, *epilogue_args) + + def compile( + self, + a_tensor, + b_tensor, + max_active_clusters, + current_stream, + *supplemental_arguments, + **compiler_options, + ): + """Compile the GEMM kernel with epilogue fusion. + + Compiles the kernel using CUTE DSL compilation, incorporating the EFC + (Epilogue Fusion Configuration) with all supplemental arguments. Returns + a callable function that accepts the same arguments for execution. + + :param a_tensor: Input tensor A + :param b_tensor: Input tensor B + :param max_active_clusters: Maximum number of active clusters + :param current_stream: CUDA stream for execution + :param supplemental_arguments: Additional arguments for the epilogue (tensors and scalars) + :param compiler_options: Keywords arguments passed to CuTe DSL compiler. "options = " for now + :return: Compiled callable function that executes the GEMM with the same signature + """ + self.efc.compile(supplemental_arguments) + compiled = cutlass.cute.compile( + self, + a_tensor, + b_tensor, + max_active_clusters, + current_stream, + self.efc.jit.pack_arguments(*supplemental_arguments), + **compiler_options, + ) + + def inject_supplemental_arguments( + a_tensor, + b_tensor, + current_stream, + *supplemental_arguments, + ): + # Run the compiled code, do not pass the constexpr parameters nor + # the compiler options. + return compiled( + a_tensor, + b_tensor, + current_stream, + self.efc.jit.pack_arguments(*supplemental_arguments), + ) + + return inject_supplemental_arguments diff --git a/examples/python/CuTeDSL/blackwell/epilogue/common_efc.py b/examples/python/CuTeDSL/blackwell/epilogue/common_efc.py new file mode 100644 index 00000000..684364be --- /dev/null +++ b/examples/python/CuTeDSL/blackwell/epilogue/common_efc.py @@ -0,0 +1,1512 @@ +# 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. + +# This not to use module annotations from future version but to change the type system to postpone the evaluation of annotations, +# about forward declaration and lazy type checking. +# See https://docs.python.org/3/library/__future__.html#future__.annotations and https://peps.python.org/pep-0563/. +from __future__ import annotations + +import dataclasses +import enum +import functools +import inspect +import logging +import os +import types +import typing + +import cutlass +import torch + +# To have some verbosity, set the CUTE_DSL_EFC_LOG_LEVEL environment variable to +# INFO or even DEBUG before launching this program. +if log_level := os.environ.get("CUTE_DSL_EFC_LOG_LEVEL", None): + logging.basicConfig(level=getattr(logging, log_level)) +logger = logging.getLogger(__name__) + + +def log(message: str): + """Helper function to log messages. Change logger.info to another level here + if needed.""" + logger.info(message) + + +""" +CUTLASS EFC Framework +""" + +# Available activation functions in the EFC Configuration class +ACTIVATION_FUNCTIONS = [ + "identity", + "relu", + "leaky_relu", + "tanh", + "sigmoid", + "silu", + "hardswish", + "gelu", +] + + +def if_debug(function): + """Execute a function if in debug mode.""" + if logger.isEnabledFor(logging.DEBUG): + function() + + +def mark_mlir(message: str): + """Put some message in MLIR output to make MLIR assembly clearer or trace execution.""" + if_debug(lambda: cutlass.cute.printf(f"mark_mlir: {message}")) + + +def trace_in_mlir(func): + """Decorator to trace function entry and exit in MLIR.""" + + @functools.wraps(func) + def wrapper(*args, **kwargs): + function_name = func.__name__ + mark_mlir(f"entering {function_name}") + result = func(*args, **kwargs) + mark_mlir(f"leaving {function_name}") + return result + + return wrapper + + +def create_named_epilogue(param_names, func): + """Create a wrapper function with specific parameter names that delegates to an implementation function. + + This function solves a common problem in the EFC (Epilogue Fusion Configuration) framework: + epilogue functions must have parameters with specific names (e.g., "alpha", "beta", "C", "D") + to match the EFC calling convention, but you may want to generate these functions + programmatically from a generic implementation. + + Instead of using string manipulation with exec() or eval() (which is insecure and breaks + tooling), this function uses Python's inspect module to create a proper function signature + that tools like debuggers, type checkers, and IDEs can understand. + + Args: + param_names: List of parameter names for the generated function + (e.g., ["alpha", "beta", "C", "x_factor"]) + func: Implementation function that accepts the same number of arguments as param_names. + The arguments will be passed in the order specified by param_names. + + Returns: + A new function with the specified parameter names that calls func with those + parameters in order. The wrapper preserves func's name and docstring, and + has a proper signature for introspection. + + Example: + # Generic implementation that doesn't care about parameter names + def compute(a, b, c): + return a + b * c + + # Create EFC-compliant function with required parameter names + epilogue = create_named_epilogue(["alpha", "X", "Y"], compute) + # Now epilogue(alpha=1, X=2, Y=3) calls compute(1, 2, 3) + # and inspect.signature(epilogue) shows the correct parameter names + + Use Case: + When programmatically generating epilogue functions with different tensor + configurations, you need each function to have the right parameter names + for the EFC framework to call them correctly with keyword arguments. + + """ + # Create Parameter objects for each parameter name, using standard Python argument binding. + parameters = [ + inspect.Parameter(name, inspect.Parameter.POSITIONAL_OR_KEYWORD) + for name in param_names + ] + + # Create a new signature with the custom parameter names + new_signature = inspect.Signature(parameters) + + # Create a wrapper function that accepts arguments according to the new signature + def wrapper(*args, **kwargs): + # Bind the provided arguments to our custom signature + bound = new_signature.bind(*args, **kwargs) + bound.apply_defaults() + + # Extract argument values in the order specified by param_names + ordered_args = [bound.arguments[name] for name in param_names] + + # Call the original function with the properly ordered arguments + return func(*ordered_args) + + # Assign the custom signature to the wrapper so introspection works correctly + wrapper.__signature__ = new_signature + wrapper.__name__ = getattr(func, "__name__", "generated_function") + wrapper.__doc__ = func.__doc__ + + return wrapper + + +class VariadicParameters: + """Minimal mixin wrapper for variadic parameters for @cute.jit/@cute.kernel + functions taking advantage that the DSL to can ingest a recursive + combination of tuples and lists.""" + + def __init__(self, efc, parameter_names): + # Add local shortcuts to the efc and gemm objects + self.efc = efc + self.gemm = efc.gemm + # Create a dataclass to have an aggregate initializer. + # Use __slots__ to avoid assigning wrong argument by error. + fields = [(name, typing.Any) for name in parameter_names] + self._parameter_class = dataclasses.make_dataclass( + "Parameter", fields, slots=True + ) + + # Add some methods to the dataclass so we can access for example arg.a + # and parameter.b also with arg["a"] or parameter["b"]. + def getitem(self, name): + """Access the dataclass attribute by name.""" + return getattr(self, name) + + self._parameter_class.__getitem__ = getitem + + def setitem(self, name, value): + """Set the dataclass attribute by name.""" + setattr(self, name, value) + + self._parameter_class.__setitem__ = setitem + + self.instantiate_args() + logger.info(f"Initial {self.arg = }") + + def pack_arguments(self, *args, **kwargs): + """Pack the arguments to pass them through a @cute.jit/@cute.kernel + call. + + If some arguments are provided, pack them, otherwise just use the + self.arg object by default. + + Return a tuple as an interface object since a @cute.jit/@cute.kernel + can ingest a recursive combination of tuples and lists.""" + if args or kwargs: + # Override the current argument object from the provided arguments, if any. + self.arg = self._parameter_class(*args, **kwargs) + # dataclasses.astuple(self.arg) breaks because it is recursive and + # applies a deepcopy incompatible with the DSL magic. Just generate 1 + # level of tuple of object references. + r = tuple(self.arg[f.name] for f in dataclasses.fields(self.arg)) + logger.info(f"pack_arguments {args = } {kwargs = } {self.arg = } {r = }") + # The DSL does not accept an empty tuple but can handle None. So + # remap to None in that case. + if not r: + return None + return r + + def unpack_parameters(self, p: typing.Tuple): + """Unpack the parameters inside a @cute.jit/@cute.kernel function. + + Assign all the self.parameter attributes.""" + # Do the opposite mapping of None to an empty tuple to have the + # parameter constructor happy. + if p is None: + p = () + # Instantiate the dataclass holding the parameters from the + # individual parameter values. + self.parameter = self._parameter_class(*p) + logger.info(f"unpack_parameters {p = } {self.parameter = }") + + def instantiate_args(self): + """Create an arg attribute from the Parameter class to be used + as an alternative way to pass the arguments instead of using an + explicit pack_arguments(). + + All the arg attributes are initialized to a noticeable name so that + any forgotten field will trigger an error.""" + + class _UnassignedArgument: + """Sentinel class to detect uninitialized arguments""" + + def __repr__(self): + return "" + + self.arg = self._parameter_class( + *([_UnassignedArgument] * len(dataclasses.fields(self._parameter_class))) + ) + + +class EFC: + """Epilogue Fusion Configuration.""" + + # Helper functions for CuTe operations + @staticmethod + def maximum(x, y): + """Element-wise maximum of 2 CuTe tensors""" + x_type = x.element_type + y_type = y.element_type + assert x_type is y_type, f"Type mismatch: x is {x_type}, y is {y_type}" + return cutlass.cute.where(x > y, x, y) + + @staticmethod + def minimum(x, y): + """Element-wise minimum of 2 CuTe tensors""" + x_type = x.element_type + y_type = y.element_type + assert x_type is y_type, f"Type mismatch: x is {x_type}, y is {y_type}" + return cutlass.cute.where(x < y, x, y) + + class JIT(VariadicParameters): + """Handle Python/@cute.jit and its boundaries with Host.""" + + # All the following customization functions should go somewhere else in + # the long term, as part of a refactoring similar to CUTLASS + # collective/main loop/epilogue... + + @trace_in_mlir + def record_tensor_dtypes(self): + """It does not seem that the tma_tensor and tma_atom carry over the + element type, so, store it here for later use.""" + self.tensor_dtype = {} + + def f(tensor_name, attributes): + tensor = self.parameter[tensor_name] + self.tensor_dtype[tensor_name] = tensor.element_type + + self.efc.foreach_tensor(f) + + @trace_in_mlir + def written_tensor_name_with_bigger_element_type(self): + """The type of the written tensor is used to compute a lot of + implementation details about tiling and so on in the kernel. + + The compilation phase has checked already there is at least 1 + written tensor name. + + Return the name of the written tensor with the biggest + element_type. + + """ + return max( + (tensor_name for tensor_name in self.efc.written_tensor_names), + key=lambda tensor_name: self.tensor_dtype[tensor_name].width, + ) + + @trace_in_mlir + def read_tensor_name_with_bigger_element_type(self): + """The type of the read tensor is used to compute a lot of + implementation details about tiling and so on in the kernel. Return the name of the read tensor with the biggest element_type, or None if there is no read tensor.""" + if self.efc.read_tensor_names: + return max( + (tensor_name for tensor_name in self.efc.read_tensor_names), + key=lambda tensor_name: self.tensor_dtype[tensor_name].width, + ) + return None + + @trace_in_mlir + def compute_stage(self): + """Get the contribution from the tensors used in the EFC to the + pipeline stage numbers.""" + self.smem_size_of_read_tensors = 0 + self.smem_size_of_written_tensors = 0 + self.tensor_dtype = {} + + def f(tensor_name, attributes): + tensor = self.parameter[tensor_name] + tensor_layout = cutlass.utils.LayoutEnum.from_tensor(tensor) + if cutlass.const_expr(self.gemm.d_layout != tensor_layout): + error_msg = ( + f"The tensor {tensor_name} has layout {tensor_layout} which is " + f"different from C/D specified layout {self.gemm.d_layout}." + ) + raise ValueError(error_msg) + + # It does not seem that the tma_tensor and tma_atom carry over + # the element type, so, store it here for later use. + self.tensor_dtype[tensor_name] = tensor.element_type + + smem_size_in_bytes_of_a_pipeline_stage = cutlass.cute.size_in_bytes( + tensor.element_type, + cutlass.utils.blackwell_helpers.make_smem_layout_epi( + tensor.element_type, self.gemm.d_layout, self.gemm.epi_tile, 1 + ), + ) + # Prepare the information to be asked soon, to recycle this + # loop. + if attributes.is_read: + self.smem_size_of_read_tensors += ( + smem_size_in_bytes_of_a_pipeline_stage + ) + if attributes.is_written: + self.smem_size_of_written_tensors += ( + smem_size_in_bytes_of_a_pipeline_stage + ) + + self.efc.foreach_tensor(f) + + @trace_in_mlir + def smem_size_in_bytes_of_read_tensors(self): + """Get the contribution in a smem pipeline stage from the tensors + loaded in the EFC.""" + logger.info(f"\t{self.smem_size_of_read_tensors = }") + return self.smem_size_of_read_tensors + + @trace_in_mlir + def smem_size_in_bytes_of_written_tensors(self): + """Get the contribution in a smem pipeline stage from the tensors + stored in the EFC.""" + logger.info(f"\t{self.smem_size_of_written_tensors = }") + return self.smem_size_of_written_tensors + + @trace_in_mlir + def smem_layout(self): + """Get the smem_layout for the tensors used in the EFC.""" + self.smem_layout_staged_read = {} + self.smem_layout_staged_written = {} + + def f(tensor_name, attributes): + tensor = self.parameter[tensor_name] + tensor_layout = cutlass.utils.LayoutEnum.from_tensor(tensor) + log(f"JIT.smem_layout {tensor_name} = {tensor!s}") + log(f"JIT.smem_layout tensor_layout[{tensor_name}] = {tensor_layout!s}") + + if attributes.is_read: + self.smem_layout_staged_read[tensor_name] = ( + cutlass.utils.blackwell_helpers.make_smem_layout_epi( + tensor.element_type, + tensor_layout, + self.gemm.epi_tile, + self.gemm.num_c_stage, + ) + ) + log(f"JIT.smem_layout read {self.gemm.num_c_stage = }") + log( + f"JIT.smem_layout read self.smem_layout_staged_read[{tensor_name}] = {self.smem_layout_staged_read[tensor_name]!s}" + ) + if attributes.is_written: + self.smem_layout_staged_written[tensor_name] = ( + cutlass.utils.blackwell_helpers.make_smem_layout_epi( + tensor.element_type, + tensor_layout, + self.gemm.epi_tile, + self.gemm.num_d_stage, + ) + ) + log(f"JIT.smem_layout written {self.gemm.num_d_stage = }") + log( + f"JIT.smem_layout written self.smem_layout_staged_written[{tensor_name}] = {self.smem_layout_staged_written[tensor_name]!s}" + ) + + self.efc.foreach_tensor(f) + + @trace_in_mlir + def create_tma_arguments(self): + """Set the TMA related arguments for the tensors used in the EFC.""" + # Make the difference for read/written to handle the case a tensor + # is both read and written. + self.total_tma_load_bytes = 0 # Used by the PipelineTmaAsync + self.tma_atom_read = {} + self.tma_tensor_read = {} + self.tma_atom_written = {} + self.tma_tensor_written = {} + + def f(tensor_name, attributes): + tensor = self.parameter[tensor_name] + + if attributes.is_read: + smem_layout = cutlass.cute.slice_( + self.smem_layout_staged_read[tensor_name], (None, None, 0) + ) + self.total_tma_load_bytes += cutlass.cute.size_in_bytes( + tensor.element_type, smem_layout + ) + ( + self.tma_atom_read[tensor_name], + self.tma_tensor_read[tensor_name], + ) = cutlass.cute.nvgpu.cpasync.make_tiled_tma_atom( + cutlass.cute.nvgpu.cpasync.CopyBulkTensorTileG2SOp(), + tensor, + smem_layout, + self.gemm.epi_tile, + ) + log( + f"JIT.tma_atom_read[{tensor_name}] = {self.tma_atom_read[tensor_name]!s}" + ) + log( + f"JIT.tma_tensor_read[{tensor_name}] = {self.tma_tensor_read[tensor_name]!s}" + ) + + if attributes.is_written: + smem_layout = cutlass.cute.slice_( + self.smem_layout_staged_written[tensor_name], (None, None, 0) + ) + ( + self.tma_atom_written[tensor_name], + self.tma_tensor_written[tensor_name], + ) = cutlass.cute.nvgpu.cpasync.make_tiled_tma_atom( + cutlass.cute.nvgpu.cpasync.CopyBulkTensorTileS2GOp(), + tensor, + smem_layout, + self.gemm.epi_tile, + ) + log( + f"JIT.tma_atom_written[{tensor_name}] = {self.tma_atom_written[tensor_name]!s}" + ) + log( + f"JIT.tma_tensor_written[{tensor_name}] = {self.tma_tensor_written[tensor_name]!s}" + ) + + self.efc.foreach_tensor(f) + + @trace_in_mlir + def create_supplemental_arguments_for_kernel(self): + """Executed before launching the @cute.kernel function to set up the + supplemental arguments to pass to the @cute.kernel function. + + In the @cute.kernel example, the parameters like `X_tma_tensor_read` + or `Y_tma_tensor_written` correspond to `mX_mnl` and `mY_mnl`.""" + argument_names = [] + + def compute_argument_names(name, attributes): + if not attributes.is_tensor: + # Just propagate the dynamic scalar with the same name. + argument_names.append(name) + else: + if attributes.is_read: + argument_names.append(f"{name}_tma_atom_read") + argument_names.append(f"{name}_tma_tensor_read") + argument_names.append(f"{name}_smem_layout_staged_read") + if attributes.is_written: + argument_names.append(f"{name}_tma_atom_written") + argument_names.append(f"{name}_tma_tensor_written") + argument_names.append(f"{name}_smem_layout_staged_written") + + self.efc.foreach_argument(compute_argument_names) + # Create the @cute.kernel-side meta-programming infrastructure + # handling also the supplemental argument handling. + self.efc.kernel = EFC.Kernel(self.efc, argument_names) + + arg = self.efc.kernel.arg + + def populate_the_kernel_arguments(name, attributes): + if not attributes.is_tensor: + # Just propagate the dynamic scalar with the same name. + arg[name] = self.parameter[name] + else: + if attributes.is_read: + arg[f"{name}_tma_atom_read"] = self.tma_atom_read[name] + arg[f"{name}_tma_tensor_read"] = self.tma_tensor_read[name] + arg[f"{name}_smem_layout_staged_read"] = ( + self.smem_layout_staged_read[name] + ) + if attributes.is_written: + arg[f"{name}_tma_atom_written"] = self.tma_atom_written[name] + arg[f"{name}_tma_tensor_written"] = self.tma_tensor_written[ + name + ] + arg[f"{name}_smem_layout_staged_written"] = ( + self.smem_layout_staged_written[name] + ) + + self.efc.foreach_argument(populate_the_kernel_arguments) + + class Kernel(VariadicParameters): + """Handle kernel part and @cute.jit/@cute.kernel boundaries.""" + + @trace_in_mlir + def prefetch_tma_descriptors(self): + """Prefetch the TMA descriptors for the tensors used in the EFC.""" + + def f(tensor_name, attributes): + if attributes.is_read: + cutlass.cute.nvgpu.cpasync.prefetch_descriptor( + self.parameter[f"{tensor_name}_tma_atom_read"] + ) + + if attributes.is_written: + cutlass.cute.nvgpu.cpasync.prefetch_descriptor( + self.parameter[f"{tensor_name}_tma_atom_written"] + ) + + self.efc.foreach_tensor(f) + + @trace_in_mlir + def allocate_smem(self): + """Allocate the shared memory for all the supplemental tensors.""" + self.smem_read = {} + self.smem_written = {} + + def f(tensor_name, attributes): + element_type = self.efc.jit.tensor_dtype[tensor_name] + if attributes.is_read: + smem_layout_staged = self.parameter[ + f"{tensor_name}_smem_layout_staged_read" + ] + self.smem_read[tensor_name] = self.gemm.smem.allocate_tensor( + element_type=element_type, + layout=smem_layout_staged.outer, + byte_alignment=self.gemm.buffer_align_bytes, + swizzle=smem_layout_staged.inner, + ) + if attributes.is_written: + smem_layout_staged = self.parameter[ + f"{tensor_name}_smem_layout_staged_written" + ] + self.smem_written[tensor_name] = self.gemm.smem.allocate_tensor( + element_type=element_type, + layout=smem_layout_staged.outer, + byte_alignment=self.gemm.buffer_align_bytes, + swizzle=smem_layout_staged.inner, + ) + + self.efc.foreach_tensor(f) + + @trace_in_mlir + def partition_global_tensors_for_tiled_mma(self): + """Partition the global supplemental tensors for TiledMMA_C/D.""" + self.tCgC_read = {} + self.tCgD_written = {} + + def f(tensor_name, attributes): + if attributes.is_read: + # (bM, bN, loopM, loopN, loopL) + gC_mnl = cutlass.cute.local_tile( + self.parameter[f"{tensor_name}_tma_tensor_read"], + cutlass.cute.slice_(self.gemm.mma_tiler, (None, None, 0)), + (None, None, None), + ) + log( + f"Kernel.partition_global_tensors_for_tiled_mma: gC_mnl[{tensor_name}] = {gC_mnl!s}" + ) + # (MMA, MMA_M, MMA_N, loopM, loopN, loopL) + self.tCgC_read[tensor_name] = self.gemm.thr_mma.partition_C(gC_mnl) + log( + f"Kernel.partition_global_tensors_for_tiled_mma: self.tCgC_read[{tensor_name}] = {self.tCgC_read[tensor_name]!s}" + ) + + if attributes.is_written: + # (bM, bN, loopM, loopN, loopL) + gD_mnl = cutlass.cute.local_tile( + self.parameter[f"{tensor_name}_tma_tensor_written"], + cutlass.cute.slice_(self.gemm.mma_tiler, (None, None, 0)), + (None, None, None), + ) + log( + f"Kernel.partition_global_tensors_for_tiled_mma: gD_mnl[{tensor_name}] = {gD_mnl!s}" + ) + # (MMA, MMA_M, MMA_N, loopM, loopN, loopL) + self.tCgD_written[tensor_name] = self.gemm.thr_mma.partition_C( + gD_mnl + ) + log( + f"Kernel.partition_global_tensors_for_tiled_mma: self.tCgD_written[{tensor_name}] = {self.tCgD_written[tensor_name]!s}" + ) + + self.efc.foreach_tensor(f) + + # The following functions are executed by the specialized warps for + # epilogue computation. + + @trace_in_mlir + def copy_and_partition_supplemental_rmem_tensors( + self, tiled_copy_t2r, tTR_rAcc, epi_tidx, epi_tile + ): + # Load tensor. + self.tiled_copy_s2r = {} + self.tSR_rC = {} + self.tSR_sC = {} + + # Store tensor. + self.tiled_copy_r2s = {} + self.tRS_rD = {} + self.tRS_sD = {} + self.bSG_sD = {} # ((ATOM_V, REST_V), EPI_M, EPI_N) + self.bSG_gD_partitioned = {} # ((ATOM_V, REST_V), EPI_M, EPI_N, RestM, RestN, RestL) + + log( + f"Kernel.copy_and_partition_supplemental_rmem_tensors: tiled_copy_t2r = {tiled_copy_t2r!s}" + ) + log( + f"Kernel.copy_and_partition_supplemental_rmem_tensors: tTR_rAcc = {tTR_rAcc!s}" + ) + log( + f"Kernel.copy_and_partition_supplemental_rmem_tensors: epi_tile = {epi_tile!s}" + ) + log( + f"Kernel.copy_and_partition_supplemental_rmem_tensors: epi_tidx = {epi_tidx!s}" + ) + + def f(tensor_name, attributes): + element_type = self.efc.jit.tensor_dtype[tensor_name] + log( + f"Kernel.copy_and_partition_supplemental_rmem_tensors: element_type[{tensor_name}] = {element_type!s}" + ) + + if attributes.is_read: + tTR_rC = cutlass.cute.make_rmem_tensor(tTR_rAcc.shape, element_type) + log( + f"Kernel.copy_and_partition_supplemental_rmem_tensors read: tTR_rC[{tensor_name}] = {tTR_rC!s}" + ) + + ( + self.tiled_copy_s2r[tensor_name], + self.tSR_rC[tensor_name], + self.tSR_sC[tensor_name], + ) = self.gemm.epilogue_smem_copy_and_partition_load( + tiled_copy_t2r, + tTR_rC, + epi_tidx, + self.smem_read[tensor_name], + ) + + log( + f"Kernel.copy_and_partition_supplemental_rmem_tensors read: self.tiled_copy_s2r[{tensor_name}] = {self.tiled_copy_s2r[tensor_name]!s}" + ) + log( + f"Kernel.copy_and_partition_supplemental_rmem_tensors read: self.tSR_rC[{tensor_name}] = {self.tSR_rC[tensor_name]!s}" + ) + log( + f"Kernel.copy_and_partition_supplemental_rmem_tensors read: self.tSR_sC[{tensor_name}] = {self.tSR_sC[tensor_name]!s}" + ) + + if attributes.is_written: + # (T2R, T2R_M, T2R_N) + tTR_rD = cutlass.cute.make_rmem_tensor(tTR_rAcc.shape, element_type) + log( + f"Kernel.copy_and_partition_supplemental_rmem_tensors written: tTR_rD[{tensor_name}] = {tTR_rD!s}" + ) + + # utils.gemm.sm100.epilogue_smem_copy_and_partition uses + # explicitly "C" as the output matrix and introspects the + # gemm object while in this kernel "C" is used for read but + # "D" is for output according to the BLAS convention. + # So construct a minimal mock-up with the required + # information. + faux_gemm = types.SimpleNamespace() + faux_gemm.c_layout = self.gemm.d_layout + faux_gemm.c_dtype = self.gemm.d_dtype + faux_gemm.acc_dtype = self.gemm.acc_dtype + ( + self.tiled_copy_r2s[tensor_name], + self.tRS_rD[tensor_name], # (R2S, R2S_M, R2S_N) + self.tRS_sD[tensor_name], # (R2S, R2S_M, R2S_N) + ) = cutlass.utils.gemm.sm100.epilogue_smem_copy_and_partition( + faux_gemm, + tiled_copy_t2r, # (EPI_TILE_M, EPI_TILE_N) + tTR_rD, + epi_tidx, + self.smem_written[tensor_name], + ) + log( + f"Kernel.copy_and_partition_supplemental_rmem_tensors written: self.smem_written[{tensor_name}] = {self.smem_written[tensor_name]!s}" + ) + log( + f"Kernel.copy_and_partition_supplemental_rmem_tensors written: self.tiled_copy_r2s[{tensor_name}] = {self.tiled_copy_r2s[tensor_name]!s}" + ) + log( + f"Kernel.copy_and_partition_supplemental_rmem_tensors written: self.tRS_rD[{tensor_name}] = {self.tRS_rD[tensor_name]!s}" + ) + log( + f"Kernel.copy_and_partition_supplemental_rmem_tensors written: self.tRS_sD[{tensor_name}] = {self.tRS_sD[tensor_name]!s}" + ) + ( + _, + self.bSG_sD[tensor_name], # ((ATOM_V, REST_V), EPI_M, EPI_N) + self.bSG_gD_partitioned[ + tensor_name + ], # ((ATOM_V, REST_V), EPI_M, EPI_N, RestM, RestN, RestL) + ) = self.gemm.epilogue_gmem_copy_and_partition( + epi_tidx, + self.parameter[f"{tensor_name}_tma_atom_written"], + self.tCgD_written[tensor_name], + epi_tile, + self.smem_written[tensor_name], + element_type, + ) + + log( + f"Kernel.copy_and_partition_supplemental_rmem_tensors written: self.parameter[{tensor_name}_tma_atom_written] = {self.parameter[f'{tensor_name}_tma_atom_written']!s}" + ) + log( + f"Kernel.copy_and_partition_supplemental_rmem_tensors written: self.tCgD_written[{tensor_name}] = {self.tCgD_written[tensor_name]!s}" + ) + log( + f"Kernel.copy_and_partition_supplemental_rmem_tensors written: self.smem_written[{tensor_name}] = {self.smem_written[tensor_name]!s}" + ) + log( + f"Kernel.copy_and_partition_supplemental_rmem_tensors written: element_type = {element_type!s}" + ) + log( + f"Kernel.copy_and_partition_supplemental_rmem_tensors written: self.bSG_sD[{tensor_name}] = {self.bSG_sD[tensor_name]!s}" + ) + log( + f"Kernel.copy_and_partition_supplemental_rmem_tensors written: self.bSG_gD_partitioned[{tensor_name}] = {self.bSG_gD_partitioned[tensor_name]!s}" + ) + + self.efc.foreach_tensor(f) + + @trace_in_mlir + def slice_written_tensors_per_mma_tile_index(self, mma_tile_coord_mnl): + """Slice the supplemental written tensors per MMA tile index.""" + self.bSG_gD = {} # ((ATOM_V, REST_V), (EPI_M, EPI_N)) + + def f(tensor_name, attributes): + # ((ATOM_V, REST_V), EPI_M, EPI_N) + bSG_gD = self.bSG_gD_partitioned[tensor_name][ + ( + None, + None, + None, + *mma_tile_coord_mnl, + ) + ] + log( + f"Kernel.slice_written_tensors_per_mma_tile_index: bSG_gD[{tensor_name}] = {bSG_gD!s}" + ) + # Group the 2 last modes so the subtile_idx loop can iterate + # through it using 1-D indexing. + # ((ATOM_V, REST_V), (EPI_M, EPI_N)) + self.bSG_gD[tensor_name] = cutlass.cute.group_modes( + bSG_gD, 1, cutlass.cute.rank(bSG_gD) + ) + log( + f"Kernel.slice_written_tensors_per_mma_tile_index: self.bSG_gD[{tensor_name}] = {self.bSG_gD[tensor_name]!s}" + ) + + self.efc.foreach_written_tensor(f) + + @trace_in_mlir + def load_tensors_from_smem_to_register(self, index): + """Load supplemental tensors from shared memory to register.""" + + def f(tensor_name, attributes): + cutlass.cute.copy( + self.tiled_copy_s2r[tensor_name], + self.tSR_sC[tensor_name][ + ( + None, + None, + None, + index, + ) + ], + self.tSR_rC[tensor_name], + ) + log( + f"Kernel.load_tensors_from_smem_to_register cutlass.cute.copy: self.tiled_copy_s2r[{tensor_name}] = {self.tiled_copy_s2r[tensor_name]!s}" + ) + log( + f"Kernel.load_tensors_from_smem_to_register cutlass.cute.copy: self.tSR_sC[{tensor_name}] = {self.tSR_sC[tensor_name]!s}" + ) + log( + f"Kernel.load_tensors_from_smem_to_register cutlass.cute.copy: self.tSR_rC[{tensor_name}] = {self.tSR_rC[tensor_name]!s}" + ) + + self.efc.foreach_read_tensor(f) + + @trace_in_mlir + def epilogue_computation(self, epilogue_context): + """Execute the EFC epilogue.""" + + epilogue_context.load = {} + epilogue_context.store = {} + + def load_setup(tensor_name, attributes): + # Retile the read subtile to fit the accumulator subtile vector + # TV layout. + epilogue_context.load[tensor_name] = ( + self.tiled_copy_r2s[self.gemm.d_name_bigger] + .retile(self.tSR_rC[tensor_name]) + .load() + ) + log( + f"Kernel.epilogue_computation load_setup: {self.tiled_copy_r2s[self.gemm.d_name_bigger] = !s}" + ) + log( + f"Kernel.epilogue_computation load_setup: self.tSR_rC[{tensor_name}] = {self.tSR_rC[tensor_name]!s}" + ) + log( + f"Kernel.epilogue_computation load_setup: self.tiled_copy_r2s[self.gemm.d_name_bigger].retile(self.tSR_rC[{tensor_name}]) = {self.tiled_copy_r2s[self.gemm.d_name_bigger].retile(self.tSR_rC[tensor_name])!s}" + ) + log( + f"Kernel.epilogue_computation load_setup: epilogue_context.load[{tensor_name}] = {epilogue_context.load[tensor_name]!s}" + ) + + self.efc.foreach_read_tensor(load_setup) + + def store_setup(tensor_name, attributes): + epilogue_context.store[tensor_name] = self.tRS_rD[tensor_name] + log( + f"Kernel.epilogue_computation store_setup: epilogue_context.store[{tensor_name}] = {epilogue_context.store[tensor_name]!s}" + ) + + self.efc.foreach_written_tensor(store_setup) + + self.efc.specialized_epilogue(EFC.Phase.ThreadOperation, epilogue_context)() + + @trace_in_mlir + def store_written_tensors_to_smem(self, d_buffer): + """Store the EFC written tensors to shared memory.""" + + def f(tensor_name, attributes): + cutlass.cute.copy( + self.tiled_copy_r2s[tensor_name], + self.tRS_rD[tensor_name], + self.tRS_sD[tensor_name][(None, None, None, d_buffer)], + ) + log( + f"Kernel.store_written_tensors_to_smem cutlass.cute.copy: self.tiled_copy_r2s[{tensor_name}] = {self.tiled_copy_r2s[tensor_name]!s}" + ) + log( + f"Kernel.store_written_tensors_to_smem cutlass.cute.copy: self.tRS_rD[{tensor_name}] = {self.tRS_rD[tensor_name]!s}" + ) + log( + f"Kernel.store_written_tensors_to_smem cutlass.cute.copy: self.tRS_sD[{tensor_name}] = {self.tRS_sD[tensor_name]!s}" + ) + log( + f"Kernel.store_written_tensors_to_smem cutlass.cute.copy: self.tRS_sD[{tensor_name}][(None, None, None, d_buffer)] = {self.tRS_sD[tensor_name][(None, None, None, d_buffer)]!s}" + ) + + self.efc.foreach_written_tensor(f) + + @trace_in_mlir + def tma_store_written_tensors_to_gmem(self, d_buffer, subtile_idx): + """Store with TMA the written EFC tensors to global memory.""" + + def f(tensor_name, attributes): + cutlass.cute.copy( + self.parameter[f"{tensor_name}_tma_atom_written"], + self.bSG_sD[tensor_name][(None, d_buffer)], + self.bSG_gD[tensor_name][(None, subtile_idx)], + ) + log( + f"Kernel.tma_store_written_tensors_to_gmem cutlass.cute.copy: self.parameter[{tensor_name}_tma_atom_written] = {self.parameter[f'{tensor_name}_tma_atom_written']!s}" + ) + log( + f"Kernel.tma_store_written_tensors_to_gmem cutlass.cute.copy: self.bSG_sD[{tensor_name}] = {self.bSG_sD[tensor_name]!s}" + ) + log( + f"Kernel.tma_store_written_tensors_to_gmem cutlass.cute.copy: self.bSG_sD[{tensor_name}][(None, d_buffer)] = {self.bSG_sD[tensor_name][(None, d_buffer)]!s}" + ) + log( + f"Kernel.tma_store_written_tensors_to_gmem cutlass.cute.copy: self.bSG_gD[{tensor_name}] = {self.bSG_gD[tensor_name]!s}" + ) + log( + f"Kernel.tma_store_written_tensors_to_gmem cutlass.cute.copy: self.bSG_gD[{tensor_name}][(None, subtile_idx)] = {self.bSG_gD[tensor_name][(None, subtile_idx)]!s}" + ) + + self.efc.foreach_written_tensor(f) + + # The following functions are executed by the specialized warp for the + # epilogue load. + + @trace_in_mlir + def create_epilogue_subtile_tensors(self, tidx, epi_tile): + """Setup the pipelines reading the EFC supplemental tensors.""" + self.bGS_sC = {} + self.bGS_gC_partitioned = {} + + def f(tensor_name, attributes): + ( + _, + self.bGS_sC[tensor_name], + self.bGS_gC_partitioned[tensor_name], + ) = self.gemm.epilogue_gmem_copy_and_partition( + tidx, + self.parameter[f"{tensor_name}_tma_atom_read"], + self.tCgC_read[tensor_name], + epi_tile, + self.smem_read[tensor_name], + self.efc.jit.tensor_dtype[tensor_name], + ) + + self.efc.foreach_read_tensor(f) + + @trace_in_mlir + def prepare_tensor_load_for_subtiles( + self, + mma_tile_coord_mnl, + ): + """Prepare the EFC tensors to be loaded by the subtiles and return the number of subtiles to compute.""" + self.bGS_gC = {} + # In the case there is no supplemental tensor to load in the + # epilogue: + self._subtile_cnt = 0 + + def f(tensor_name, attributes): + self.bGS_gC[tensor_name] = self.bGS_gC_partitioned[tensor_name][ + ( + None, + None, + None, + *mma_tile_coord_mnl, + ) + ] + self.bGS_gC[tensor_name] = cutlass.cute.group_modes( + self.bGS_gC[tensor_name], + 1, + cutlass.cute.rank(self.bGS_gC[tensor_name]), + ) + st_cnt = cutlass.cute.size(self.bGS_gC[tensor_name].shape, mode=[1]) + if self._subtile_cnt == 0: + # Keep the first loaded tensor as a reference. + self._subtile_cnt = st_cnt + if st_cnt != self._subtile_cnt: + raise NotImplementedError( + f"Subtile count mismatch: tensor '{self.efc.read_tensor_names[0]}' has {self._subtile_cnt} subtiles, " + f"but tensor '{tensor_name}' has {st_cnt} subtiles. All tensors must have the same subtile count." + ) + + self.efc.foreach_read_tensor(f) + + return self._subtile_cnt + + @trace_in_mlir + def load_tensor_subtiles( + self, subtile_idx, c_pipeline, c_pipeline_producer_state + ): + """Load the subtiles of the EFC tensors.""" + + def f(tensor_name, attributes): + # Load supplemental tensor from global memory to shared memory. + cutlass.cute.copy( + self.parameter[f"{tensor_name}_tma_atom_read"], + self.bGS_gC[tensor_name][(None, subtile_idx)], + self.bGS_sC[tensor_name][(None, c_pipeline_producer_state.index)], + tma_bar_ptr=c_pipeline.producer_get_barrier( + c_pipeline_producer_state + ), + ) + + self.efc.foreach_read_tensor(f) + + class Phase(enum.Enum): + ParameterAnalysis = enum.auto() + """Epilogue function during analysis of its parameters.""" + + ThreadOperation = enum.auto() + """Epilogue function used for computation.""" + + PyTorchEvaluation = enum.auto() + """Epilogue function used for verification on CPU with PyTorch.""" + + class Tensor: + """A proxy object to be used as an argument to introspect or execute the + epilogue configuration function in a given phase.""" + + @dataclasses.dataclass + class ParameterAttributes: + """Store some characteristics of the epilogue parameters""" + + is_tensor: bool # Tensor or scalar + is_read: bool = False + is_written: bool = False + + def __init__( + self, + phase: typing.ForwardRef("EFC.Phase"), + name: str, + efc: EFC, + configuration, + ): + self.phase = phase + self.name = name + self.efc = efc + self.configuration = configuration + self.attributes: EFC.Tensor.ParameterAttributes = efc.parameter_attributes[ + name + ] + logger.info(f"Tensor {self.name = }") + + def load(self): + """""" + match self.phase: + case EFC.Phase.ParameterAnalysis: + # Record that the tensor is read: + self.attributes.is_read = True + # Some value to have expression evaluation happy + return 1 + + case EFC.Phase.ThreadOperation: + # arg[0] is the epilogue_context from epilogue_computation(). + return ( + self.configuration.args[0] + .load[self.name] + .to(self.efc.gemm.epi_dtype) + ) + + case EFC.Phase.PyTorchEvaluation: + # args[1] is VariadicParameters constructed in + # evaluate_on_cpu(). Use .arg and not .parameter since it is + # not used actually to handle variadic parameter passing + # here. Just return the PyTorch tensor. + # TODO: Need to map to matching cutlass type. + return self.configuration.args[1].arg[self.name] + + case _: + raise NotImplementedError( + f"load({self.name}) not implemented for phase {self.phase}" + ) + + def store(self, value): + """""" + match self.phase: + case EFC.Phase.ParameterAnalysis: + # Record that the tensor is written: + self.attributes.is_written = True + + case EFC.Phase.ThreadOperation: + # arg[0] is the epilogue_context from epilogue_computation(). + tRS_rD = self.configuration.args[0].store[self.name] + tRS_rD.store(value.to(self.efc.jit.tensor_dtype[self.name])) + + case EFC.Phase.PyTorchEvaluation: + # args[1] is VariadicParameters constructed in + # evaluate_on_cpu(). Use .arg and not .parameter since it is + # not used actually to handle variadic parameter passing + # here. Assign the PyTorch tensor target with the given + # value. + self.configuration.args[1].arg[self.name].copy_(value) + + case _: + raise NotImplementedError( + f"store({self.name}) not implemented for phase {self.phase}" + ) + + class Configuration: + """Specialize the epilogue provided by the user to be called in the + compute kernel customization point at a given phase.""" + + def __init__(self, efc: EFC, phase: EFC.Phase, *args): + """""" + self.efc = efc + self.phase = phase + # args[0] is the epilogue_context from the kernel for EFC.Phase.ThreadOperation. + self.args = args + self.arguments = [ + self._argument(name) for name in efc.epilogue_parameter_names + ] + + def _argument(self, name): + """Generate the argument used by the specialized epilogue with the + given name""" + if self.efc.parameter_attributes[name].is_tensor: + # Delegate the phase-related behavior to the Tensor object + # itself. + return EFC.Tensor(self.phase, name, self.efc, self) + # Otherwise, we have a dynamic scalar parameter. + match self.phase: + case EFC.Phase.ParameterAnalysis: + # Use some dummy value during introspection phase. + return cutlass.Float32(42).to(self.efc.gemm.epi_dtype) + + case EFC.Phase.ThreadOperation: + # TODO: Need to map to matching cutlass type. + # Return directly the real kernel parameter with the same name. + return cutlass.Float32(self.efc.kernel.parameter[name]).to( + self.efc.gemm.epi_dtype + ) + + case EFC.Phase.PyTorchEvaluation: + # args[1] is VariadicParameters constructed in + # evaluate_on_cpu(). Use .arg and not .parameter since it is + # not used actually to handle variadic parameter passing + # here. + # TODO: Need to map to matching cutlass type. + return self.args[1].arg[name] + + case _: + raise NotImplementedError( + f"argument({name}) not implemented for phase {self.phase}" + ) + + def __call__(self): + """Execute the epilogue provided by the end-user with some specific + arguments crafted for the current phase. + + Pass self as an argument, to be seen as `efc_config`, a way to + access the EFC instance and its properties.""" + return self.efc.epilogue_function_configuration(self, *self.arguments) + + def accum(self): + """Provide the accumulator value to the user.""" + match self.phase: + case EFC.Phase.ParameterAnalysis: + # The answer to anything. + return cutlass.Float32(42) + + case EFC.Phase.ThreadOperation: + # args[0] is epilogue_context passed to + # efc.kernel.epilogue_computation(). + return self.args[0].acc_vec + + case EFC.Phase.PyTorchEvaluation: + # Return matrix_multiplication_ref from evaluate_on_cpu(). + return self.args[0] + + case _: + raise NotImplementedError( + f"accum() not implemented for phase {self.phase}" + ) + + # Some helper functions for common operations. + + def maximum(self, x, y): + """Element-wise maximum of 2 tensors""" + match self.phase: + case EFC.Phase.ParameterAnalysis: + return 1 + case EFC.Phase.ThreadOperation: + return EFC.maximum(x, y) + case EFC.Phase.PyTorchEvaluation: + return torch.maximum(x, y) + case _: + raise NotImplementedError( + f"maximum() not implemented for phase {self.phase}" + ) + + def minimum(self, x, y): + """Element-wise minimum of 2 tensors""" + match self.phase: + case EFC.Phase.ParameterAnalysis: + return 1 + case EFC.Phase.ThreadOperation: + return EFC.minimum(x, y) + case EFC.Phase.PyTorchEvaluation: + return torch.minimum(x, y) + case _: + raise NotImplementedError( + f"minimum() not implemented for phase {self.phase}" + ) + + # Define some activation functions inspired by: + # - cutlass/python/cutlass_cppgen/epilogue/epilogue.py + # - cutlass/python/cutlass_cppgen/backend/epilogue.py + + def identity(self, x): + """Identity activation function: f(x) = x""" + match self.phase: + case EFC.Phase.ParameterAnalysis: + return 1 + case EFC.Phase.ThreadOperation: + return x + case EFC.Phase.PyTorchEvaluation: + return x + case _: + raise NotImplementedError( + f"identity() not implemented for phase {self.phase}" + ) + + def relu(self, x): + """ReLU activation function: f(x) = maximum(0, x)""" + match self.phase: + case EFC.Phase.ParameterAnalysis: + return 1 + case EFC.Phase.ThreadOperation: + return EFC.maximum(x, self.full_like(x, 0)) + case EFC.Phase.PyTorchEvaluation: + return torch.nn.functional.relu(x) + case _: + raise NotImplementedError( + f"relu() not implemented for phase {self.phase}" + ) + + def leaky_relu(self, x, negative_slope=0.01): + """Leaky ReLU activation function: f(x) = maximum(0, x) + negative_slope * minimum(0, x)""" + match self.phase: + case EFC.Phase.ParameterAnalysis: + return 1 + case EFC.Phase.ThreadOperation: + # Use self.full_like to have all the computation done with + # same type as x element type. + zero = self.full_like(x, 0) + return EFC.maximum(x, zero) + EFC.minimum(x, zero) * self.full_like( + x, negative_slope + ) + case EFC.Phase.PyTorchEvaluation: + return torch.nn.functional.leaky_relu(x, negative_slope) + case _: + raise NotImplementedError( + f"leaky_relu() not implemented for phase {self.phase}" + ) + + def tanh(self, x): + """Hyperbolic tangent activation function""" + match self.phase: + case EFC.Phase.ParameterAnalysis: + return 1 + case EFC.Phase.ThreadOperation: + return cutlass.cute.tanh(x) + case EFC.Phase.PyTorchEvaluation: + return torch.tanh(x) + case _: + raise NotImplementedError( + f"tanh() not implemented for phase {self.phase}" + ) + + def sigmoid(self, x): + """Sigmoid activation function: f(x) = 1 / (1 + exp(-x))""" + match self.phase: + case EFC.Phase.ParameterAnalysis: + return 1 + case EFC.Phase.ThreadOperation: + # Use self.full_like to have all the computation done with + # same type as x element type. + # sigmoid(x) = 1 / (1 + exp(-x)) + return self.full_like(x, 1) / ( + self.full_like(x, 1) + cutlass.cute.exp(-x) + ) + case EFC.Phase.PyTorchEvaluation: + return torch.sigmoid(x) + case _: + raise NotImplementedError( + f"sigmoid() not implemented for phase {self.phase}" + ) + + def silu(self, x): + """SiLU (Swish) activation function: f(x) = x * sigmoid(x)""" + match self.phase: + case EFC.Phase.ParameterAnalysis: + return 1 + case EFC.Phase.ThreadOperation: + # silu(x) = x * sigmoid(x) + return x * self.sigmoid(x) + case EFC.Phase.PyTorchEvaluation: + return torch.nn.functional.silu(x) + case _: + raise NotImplementedError( + f"silu() not implemented for phase {self.phase}" + ) + + def hardswish(self, x): + """Hard Swish activation function: f(x) = x * relu6(x + 3) / 6""" + match self.phase: + case EFC.Phase.ParameterAnalysis: + return 1 + case EFC.Phase.ThreadOperation: + # Use self.full_like to have all the computation done with + # same type as x element type. + # hardswish(x) = x * minimum(maximum(x + 3, 0), 6) / 6 + relu6 = EFC.minimum( + EFC.maximum(x + self.full_like(x, 3), self.full_like(x, 0)), + self.full_like(x, 6), + ) + return x * relu6 / self.full_like(x, 6) + case EFC.Phase.PyTorchEvaluation: + return torch.nn.functional.hardswish(x) + case _: + raise NotImplementedError( + f"hardswish() not implemented for phase {self.phase}" + ) + + def gelu(self, x): + """GELU (Gaussian Error Linear Unit) activation function.""" + match self.phase: + case EFC.Phase.ParameterAnalysis: + return 1 + case EFC.Phase.ThreadOperation: + # Use self.full_like to have all the computation done with + # same type as x element type. + # GELU approximation: 0.5 * x * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x^3))) + # Using a simpler approximation for CuTe + sqrt_2_over_pi = self.full_like(x, 0.7978845608028654) + return ( + self.full_like(x, 0.5) + * x + * ( + self.full_like(x, 1) + + cutlass.cute.tanh( + sqrt_2_over_pi + * (x + self.full_like(x, 0.044715) * x * x * x) + ) + ) + ) + case EFC.Phase.PyTorchEvaluation: + return torch.nn.functional.gelu(x) + case _: + raise NotImplementedError( + f"gelu() not implemented for phase {self.phase}" + ) + + def __getattr__(self, name): + """Called when an attribute or method is not found. + + Hijack this mechanism to dispatch/emulate functions like + cute.full_like() or torch.full_like() provided inside the epilogue + function as self.full_like(). + + This is required since the epilogue is used not only in a @cute.jit + or @cute.kernel but also executed in a normal context for analyzing + the epilogue content and even run in emulation with frameworks like + PyTorch.""" + + def chameleon(self, *args, **kwargs): + """The great impostor method. + + TODO: add some level of configuration to tweak the CuTe/Python + name mapping, handle some specific default values for some + parameters...""" + + match self.phase: + case EFC.Phase.ParameterAnalysis: + # Just return a value to go on with the fake evaluation, in + # the case the function is expected to return a result. It + # will be ignored anyway in the opposite case. + return 1 + + case EFC.Phase.ThreadOperation: + # In the @cute.kernel context, just use the normal CuTe + # implementation. + return getattr(cutlass.cute, name)(*args, **kwargs) + + case EFC.Phase.PyTorchEvaluation: + # In the PyTorch context, call the equivalent function + # with the same name. + return getattr(torch, name)(*args, **kwargs) + + case _: + raise NotImplementedError( + f"self.{name} not implemented for phase {self.phase}" + ) + + # Update the function name to match the requested attribute name. + chameleon.__name__ = name + # Return chameleon blessed as a bound method of self. + return types.MethodType(chameleon, self) + + def __init__( + self, + gemm, + epilogue_function_configuration, + ): + """Construct an EFC instance.""" + self.gemm = gemm + self.epilogue_function_configuration = epilogue_function_configuration + self.analyze_epilogue(epilogue_function_configuration) + + def analyze_epilogue(self, epilogue_function_configuration): + """Analyze the epilogue configuration function to extract its parameter + names.""" + sig = inspect.signature(epilogue_function_configuration) + names = [name for name in sig.parameters.keys()] + # Impose to have the first parameter named "efc_config". This is very + # intrusive but at the same time some people got confused when they + # forgot this parameter. + if names[0] != "efc_config": + raise RuntimeError( + "The epilogue configuration function must take efc_config as an argument" + ) + + # Keep all the argument names but the first "efc_config" one. + self.epilogue_parameter_names = names[1:] + logger.info(f"{self.epilogue_parameter_names = }") + + def compile(self, supplemental_arguments): + """Compile with all the arguments to know the types during compilation + while hiding the epilogue detail1s.""" + assert len(supplemental_arguments) == len(self.epilogue_parameter_names) + # Update the active epilogue instance to use the new Parameter class + self.analyze_epilogue_with_arguments(supplemental_arguments) + # Create the metaprogramming objects for the @cute.jit and @cute.kernel + # parts. For now just forward all the parameters as is. + self.jit = EFC.JIT(self, self.epilogue_parameter_names) + if not self.written_tensor_names: + raise NotImplementedError( + "The epilogue requires at least one written tensor to do something useful." + ) + + def analyze_epilogue_with_arguments(self, supplemental_arguments): + self.parameter_attributes = {} + logger.info(f"{self.analyze_epilogue_with_arguments}:") + for name, a in zip(self.epilogue_parameter_names, supplemental_arguments): + logger.info(f"{name = } {a = }, {type(a) = }") + self.parameter_attributes[name] = EFC.Tensor.ParameterAttributes( + is_tensor=isinstance(a, cutlass.cute.Tensor) + ) + + # Evaluate the epilogue function for parameter analysis + self.specialized_epilogue(EFC.Phase.ParameterAnalysis)() + logger.info(f"\t{self.parameter_attributes = }") + # Keep track of all the epilogue tensor use cases per name: + self.used_tensor_names = [] + self.read_tensor_names = [] + self.written_tensor_names = [] + for name in self.epilogue_parameter_names: + q = self.parameter_attributes[name] + if not q.is_tensor: + continue + if q.is_read or q.is_written: + self.used_tensor_names.append(name) + if q.is_read: + self.read_tensor_names.append(name) + if q.is_written: + self.written_tensor_names.append(name) + logger.info( + f"\t{self.used_tensor_names = }\n\t{self.read_tensor_names = }\n\t{self.written_tensor_names = }" + ) + + def specialized_epilogue(self, phase: typing.ForwardRef("EFC.Phase"), *args): + """Construct a configuration of the epilogue specialized for a given + phase. The arguments are opaque and depend on the actual phase use.""" + return EFC.Configuration(self, phase, *args) + + def foreach_argument(self, function): + """Execute the given function for each supplemental argument of the epilogue.""" + for name in self.epilogue_parameter_names: + attributes = self.parameter_attributes[name] + function(name, attributes) + + def foreach_tensor(self, function): + """Execute the given function for each supplemental tensor.""" + for tensor_name in self.used_tensor_names: + attributes = self.parameter_attributes[tensor_name] + function(tensor_name, attributes) + + def foreach_read_tensor(self, function): + """Execute the given function for each supplemental read tensor.""" + + for tensor_name in self.read_tensor_names: + attributes = self.parameter_attributes[tensor_name] + function(tensor_name, attributes) + + def foreach_written_tensor(self, function): + """Execute the given function for each supplemental written tensor.""" + + for tensor_name in self.written_tensor_names: + attributes = self.parameter_attributes[tensor_name] + function(tensor_name, attributes) + + def evaluate_on_cpu(self, matrix_multiplication_ref, *args): + """Evaluate the epilogue fusion configuration function on CPU for + validation using the precomputed matrix multiplication result. + + Use PyTorch for now but could be whatever.""" + # Recycle the VariadicParameters class to map the arguments according to + # their names: + epilogue_args = VariadicParameters(self, self.epilogue_parameter_names) + epilogue_args.pack_arguments(*args) + # Evaluate the epilogue with PyTorch. The tensor arguments which are + # stored are also evaluated and this is how some results are returned. + self.specialized_epilogue( + EFC.Phase.PyTorchEvaluation, + matrix_multiplication_ref, + epilogue_args, + )() diff --git a/examples/python/CuTeDSL/blackwell/epilogue/custom_epilogue_dense_gemm.py b/examples/python/CuTeDSL/blackwell/epilogue/custom_epilogue_dense_gemm.py new file mode 100644 index 00000000..954cf480 --- /dev/null +++ b/examples/python/CuTeDSL/blackwell/epilogue/custom_epilogue_dense_gemm.py @@ -0,0 +1,625 @@ +# 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. + +import traceback +import typing + +import cuda.bindings.driver as cuda + +# Required for pre-Python 3.12 instead of typing.override. +from typing_extensions import override +import torch + +import cutlass +import cutlass.cute.testing as testing +import cutlass.torch as cutlass_torch + +from common_dense_gemm_efc import DenseGemmEFC + +""" +A high-performance persistent batched dense GEMM with custom epilogue fusion for the NVIDIA Blackwell SM100 architecture +using CUTE DSL and Epilogue Fusion Configuration (EFC). + +This example demonstrates a GEMM with a custom fused epilogue that performs: + Y = A * B (accumulator stored to Y) + D = (A * B) * alpha + C * beta + X * x_factor + +Tensor dimensions: +- Matrix A is MxKxL, L is batch dimension, A can be row-major("K") or column-major("M") +- Matrix B is NxKxL, L is batch dimension, B can be row-major("N") or column-major("K") +- Matrix C is MxNxL (read-only input), C can be row-major("N") or column-major("M") +- Matrix D is MxNxL (output), D can be row-major("N") or column-major("M") +- Matrix X is MxNxL (read-only input), same layout as C/D +- Matrix Y is MxNxL (output), same layout as C/D +- alpha, beta, and x_factor are scalar scale factors + +This GEMM kernel supports the following features: + - Utilizes Tensor Memory Access (TMA) for efficient memory operations + - Utilizes Blackwell's tcgen05.mma for matrix multiply-accumulate (MMA) operations (including 2cta mma instructions) + - Implements TMA multicast with cluster to reduce L2 memory traffic + - Supports persistent tile scheduling to better overlap memory load/store with mma between tiles + - Supports warp specialization to avoid explicit pipelining between mainloop load and mma + - Uses Epilogue Fusion Configuration (EFC) to define custom epilogue operations + +This GEMM works as follows: +1. DMA warp: Load A and B matrices from global memory (GMEM) to shared memory (SMEM) using TMA operations. +2. MMA warp: Perform matrix multiply-accumulate (MMA) operations using tcgen05.mma instruction. +3. EPILOGUE warp (defined via EFC): + - Load completed accumulator from tensor memory (TMEM) to registers (RMEM) using tcgen05.ld. + - Load C and X matrices from global memory (GMEM) to shared memory (SMEM) using TMA, then to registers (RMEM). + - Compute Y = accumulator (copy of A*B result) + - Compute D = accumulator * alpha + C * beta + X * x_factor + - Type convert D and Y matrices to output types. + - Store D and Y matrices from registers (RMEM) to shared memory (SMEM) to global memory (GMEM) with TMA operations + +SM100 tcgen05.mma instructions operate as follows: +- Read matrix A from SMEM +- Read matrix B from SMEM +- Write accumulator to TMEM +The accumulator in TMEM must then be loaded to registers before writing back to GMEM. + +Example usage: + +.. code-block:: bash + + python custom_epilogue_dense_gemm.py \ + --ab_dtype Float16 --c_dtype Float16 --d_dtype Float16 --acc_dtype Float32 --epi_dtype Float32 \ + --x_dtype Float16 --y_dtype Float16 \ + --mma_tiler_mn 128,128 --cluster_shape_mn 2,1 \ + --mnkl 8192,8192,8192,1 \ + --use_2cta_instrs --alpha 2.0 --beta 1.0 --x_factor 3.0 + +To collect performance with NCU profiler: + +.. code-block:: bash + + ncu python custom_epilogue_dense_gemm.py \ + --ab_dtype Float16 --c_dtype Float16 --d_dtype Float16 --acc_dtype Float32 --epi_dtype Float32 \ + --x_dtype Float16 --y_dtype Float16 \ + --mma_tiler_mn 128,128 --cluster_shape_mn 2,1 \ + --mnkl 8192,8192,8192,1 \ + --use_2cta_instrs --alpha 2.0 --beta 1.0 --x_factor 3.0 \ + --warmup_iterations 1 --iterations 10 --skip_ref_check + +Constraints: +* Supported input data types: fp16, bf16, tf32, int8, uint8, fp8 (e4m3fn, e5m2) +* A/B tensors must have the same data type +* C/D/X/Y tensors must have the same major order +* MMA tiler M must be 64/128 (use_2cta_instrs=False) or 128/256 (use_2cta_instrs=True) +* MMA tiler N must be 32-256, step 32 +* Cluster shape M/N must be positive and power of 2, total cluster size <= 16 +* Cluster shape M must be multiple of 2 if use_2cta_instrs=True +* The contiguous dimension of all tensors must be at least 16 bytes aligned, + i.e, number of elements is a multiple of 4, 8, and 16 for TFloat32, + Float16/BFloat16, and Int8/Uint8/Float8, respectively. +* OOB tiles are not allowed when TMA store is disabled +""" + + +class DenseGemmAlphaBeta(DenseGemmEFC): + """Implements batched GEMM with custom epilogue fusion using EFC. + + This class extends DenseGemmEFC to provide a fused epilogue that: + - Reads from input tensors C and X + - Writes to output tensors D and Y + - Performs: Y = A*B and D = (A*B) * alpha + C * beta + X * x_factor + + The class provides CLI argument parsing and tensor creation for the + specific epilogue configuration with C, D, X, Y tensors and alpha, + beta, x_factor scalar parameters. + """ + + class CLIParser(DenseGemmEFC.CLIParser): + @override + def more_parsing(self): + self.parser.add_argument( + "--alpha", type=float, default=1.0, help="alpha scale factor" + ) + self.parser.add_argument( + "--beta", type=float, default=0.0, help="beta scale factor" + ) + self.parser.add_argument( + "--c_dtype", type=cutlass.dtype, default=cutlass.Float32 + ) + self.parser.add_argument( + "--d_dtype", type=cutlass.dtype, default=cutlass.Float32 + ) + self.parser.add_argument( + "--x_dtype", type=cutlass.dtype, default=cutlass.Float32 + ) + self.parser.add_argument( + "--x_factor", type=float, default=3.0, help="x_factor scale factor" + ) + self.parser.add_argument( + "--y_dtype", type=cutlass.dtype, default=cutlass.Float32 + ) + + @override + def create_arguments( + self, + l, + m, + n, + k, + a_major, + b_major, + cd_major, + ab_dtype, + # For the supplemental tensors. + c_dtype, + d_dtype, + x_dtype, + y_dtype, + ): + """Create arguments for GEMM operations with epilogue tensors. + + Creates tensors for A, B (from parent class) and epilogue-specific + tensors C, D, X, Y with appropriate data types and layouts. + + :return: Tuple of (a_tensor, b_tensor, a_torch_cpu, b_torch_cpu, + c_tensor, c_torch_cpu, c_torch_gpu, + d_tensor, d_torch_cpu, d_torch_gpu, + x_tensor, x_torch_cpu, x_torch_gpu, + y_tensor, y_torch_cpu, y_torch_gpu) + """ + # Get standard arguments from parent class + std_args = super().create_arguments( + l, m, n, k, a_major, b_major, cd_major, ab_dtype + ) + + # Add the auxiliary accumulator tensors + c_torch_cpu = cutlass_torch.matrix(l, m, n, cd_major == "m", c_dtype) + c_tensor, c_torch_gpu = cutlass_torch.cute_tensor_like( + c_torch_cpu, c_dtype, is_dynamic_layout=True, assumed_align=16 + ) + + d_torch_cpu = cutlass_torch.matrix(l, m, n, cd_major == "m", d_dtype) + d_tensor, d_torch_gpu = cutlass_torch.cute_tensor_like( + d_torch_cpu, d_dtype, is_dynamic_layout=True, assumed_align=16 + ) + + x_torch_cpu = cutlass_torch.matrix(l, m, n, cd_major == "m", x_dtype) + x_tensor, x_torch_gpu = cutlass_torch.cute_tensor_like( + x_torch_cpu, x_dtype, is_dynamic_layout=True, assumed_align=16 + ) + + y_torch_cpu = cutlass_torch.matrix(l, m, n, cd_major == "m", y_dtype) + y_tensor, y_torch_gpu = cutlass_torch.cute_tensor_like( + y_torch_cpu, y_dtype, is_dynamic_layout=True, assumed_align=16 + ) + + return ( + *std_args, + c_tensor, + c_torch_cpu, + c_torch_gpu, + d_tensor, + d_torch_cpu, + d_torch_gpu, + x_tensor, + x_torch_cpu, + x_torch_gpu, + y_tensor, + y_torch_cpu, + y_torch_gpu, + ) + + def compare( + self, + a_torch_cpu, + b_torch_cpu, + epi_dtype, + tolerance, + # For the tensor check. + c_torch_gpu, + d_torch_gpu, + x_torch_gpu, + y_torch_gpu, + # The EFC epilogue arguments. + c_torch_cpu, + d_torch_cpu, + alpha, + beta, + x_torch_cpu, + x_factor, + y_torch_cpu, + ): + """Compare GPU results against CPU reference implementation. + + Evaluates the epilogue computation on CPU and validates that: + - Output tensor D matches CPU computation + - Output tensor Y matches CPU computation + - Input tensors C and X remain unchanged (read-only) + + :param a_torch_cpu: Input matrix A on CPU + :param b_torch_cpu: Input matrix B on CPU + :param epi_dtype: Data type for epilogue computation + :param tolerance: Tolerance for numerical comparison + :param c_torch_gpu: Input matrix C on GPU (to verify unchanged) + :param d_torch_gpu: Output matrix D on GPU (to compare) + :param x_torch_gpu: Input matrix X on GPU (to verify unchanged) + :param y_torch_gpu: Output matrix Y on GPU (to compare) + :param c_torch_cpu: Input matrix C on CPU + :param d_torch_cpu: Output matrix D on CPU (reference) + :param alpha: Scale factor for accumulator + :param beta: Scale factor for C + :param x_torch_cpu: Input matrix X on CPU + :param x_factor: Scale factor for X + :param y_torch_cpu: Output matrix Y on CPU (reference) + """ + # Compute reference result + self.evaluate_on_cpu( + a_torch_cpu, + b_torch_cpu, + epi_dtype, + c_torch_cpu, + d_torch_cpu, + alpha, + beta, + x_torch_cpu, + x_factor, + y_torch_cpu, + ) + # Assert close results. + torch.testing.assert_close( + d_torch_gpu.cpu(), d_torch_cpu, atol=tolerance, rtol=1e-05 + ) + torch.testing.assert_close( + y_torch_gpu.cpu(), y_torch_cpu, atol=tolerance, rtol=1e-05 + ) + # Assert that the read tensors has not been changed. + torch.testing.assert_close( + c_torch_gpu.cpu(), c_torch_cpu, atol=tolerance, rtol=1e-05 + ) + torch.testing.assert_close( + x_torch_gpu.cpu(), x_torch_cpu, atol=tolerance, rtol=1e-05 + ) + + @staticmethod + def format_as_cli_args( + mnkl: typing.Tuple[int, int, int, int], + ab_dtype: typing.Type[cutlass.Numeric], + acc_dtype: typing.Type[cutlass.Numeric], + epi_dtype: typing.Type[cutlass.Numeric], + a_major: str, + b_major: str, + cd_major: str, + mma_tiler_mn: typing.Tuple[int, int], + cluster_shape_mn: typing.Tuple[int, int], + use_2cta_instrs: bool, + c_dtype: typing.Type[cutlass.Numeric], + d_dtype: typing.Type[cutlass.Numeric], + alpha: float, + beta: float, + x_dtype: typing.Type[cutlass.Numeric], + x_factor: float, + y_dtype: typing.Type[cutlass.Numeric], + tolerance: float, + ) -> str: + """Format test parameters as CLI arguments for custom_epilogue_dense_gemm.py + + Formats all test parameters into a CLI command that can be directly + copy-pasted to reproduce the test case. Includes base parameters from + DenseGemmEFC and epilogue-specific parameters (c_dtype, d_dtype, x_dtype, + y_dtype, x_factor). + + :return: Formatted CLI command string + """ + # Get base command from parent class + base_cmd = DenseGemmEFC.format_as_cli_args( + "custom_epilogue_dense_gemm.py", + mnkl, + ab_dtype, + acc_dtype, + epi_dtype, + a_major, + b_major, + cd_major, + mma_tiler_mn, + cluster_shape_mn, + use_2cta_instrs, + tolerance, + ) + + # Add epilogue-specific parameters + specific_args = ( + f" --alpha {alpha} " + f"--beta {beta} " + f"--c_dtype {DenseGemmEFC.dtype_name(c_dtype)} " + f"--d_dtype {DenseGemmEFC.dtype_name(d_dtype)} " + f"--x_dtype {DenseGemmEFC.dtype_name(x_dtype)} " + f"--y_dtype {DenseGemmEFC.dtype_name(y_dtype)} " + f"--x_factor {x_factor}" + ) + + return base_cmd + specific_args + + +def run( + mnkl: typing.Tuple[int, int, int, int], + ab_dtype: typing.Type[cutlass.Numeric], + acc_dtype: typing.Type[cutlass.Numeric], + epi_dtype: typing.Type[cutlass.Numeric], + a_major: str, + b_major: str, + cd_major: str, + # Epilogue EFC arguments. + c_dtype: typing.Type[cutlass.Numeric], + d_dtype: typing.Type[cutlass.Numeric], + alpha: float, + beta: float, + x_dtype: typing.Type[cutlass.Numeric], + x_factor: float, + y_dtype: typing.Type[cutlass.Numeric], + # Common arguments. + mma_tiler_mn: typing.Tuple[int, int], + cluster_shape_mn: typing.Tuple[int, int], + use_2cta_instrs: bool, + tolerance: float, + warmup_iterations: int = 0, + iterations: int = 1, + skip_ref_check: bool = False, +): + """Execute batched GEMM with custom epilogue fusion. + + Performs: + Y = A * B + D = (A * B) * alpha + C * beta + X * x_factor + + :param mnkl: Matrix dimensions (M, N, K, L) where L is batch dimension + :param ab_dtype: Data type for input matrices A and B + :param acc_dtype: Data type for accumulator + :param epi_dtype: Data type for epilogue computation + :param a_major: Major order for A matrix ('k' or 'm') + :param b_major: Major order for B matrix ('k' or 'n') + :param cd_major: Major order for C/D/X/Y matrices ('n' or 'm') + :param c_dtype: Data type for input matrix C + :param d_dtype: Data type for output matrix D + :param alpha: Scale factor for accumulator in D computation + :param beta: Scale factor for C in D computation + :param x_dtype: Data type for input matrix X + :param x_factor: Scale factor for X in D computation + :param y_dtype: Data type for output matrix Y + :param mma_tiler_mn: MMA tiler dimensions (M, N) + :param cluster_shape_mn: Cluster shape (M, N) + :param use_2cta_instrs: Whether to use 2CTA instructions + :param tolerance: Tolerance for validation + :param warmup_iterations: Number of warmup iterations + :param iterations: Number of iterations to run + :param skip_ref_check: Skip reference checking + """ + print("Running Blackwell Persistent Dense GEMM test with:") + print(f"mnkl: {mnkl}") + print(f"AB dtype: {ab_dtype}, Acc dtype: {acc_dtype}, Epi dtype: {epi_dtype}") + print( + f"Matrix majors - A: {a_major}, B: {b_major}, loaded: {cd_major}, stored: {cd_major}" + ) + print(f"Mma Tiler (M, N): {mma_tiler_mn}, Cluster Shape (M, N): {cluster_shape_mn}") + print(f"2CTA MMA instructions: {'True' if use_2cta_instrs else 'False'}") + print(f"Tolerance: {tolerance}") + print(f"Warmup iterations: {warmup_iterations}") + print(f"Iterations: {iterations}") + print(f"Skip reference checking: {skip_ref_check}") + print("Epilogue:") + print(f"\t{c_dtype = !s}, {d_dtype = !s}") + print(f"\t{alpha = }, {beta = }") + print(f"\t{x_dtype = !s}, {x_factor = }") + print(f"\t{y_dtype = !s}") + + # Unpack parameters + m, n, k, l = mnkl + + if not torch.cuda.is_available(): + raise RuntimeError("GPU is required to run this example!") + + # Get current CUDA stream from PyTorch + torch_stream = torch.cuda.current_stream() + # Get the raw stream pointer as a CUstream + current_stream = cuda.CUstream(torch_stream.cuda_stream) + + # The order of the parameters here is defining the one to be used in all the + # other API calls. The epilogue function does not return anything and at + # least one tensor .store() is required to have a useful computation. + # efc_config exposes many features to the programmer, like activation + # functions or accessing some implementation details. See EFC.Configuration. + def epilogue(efc_config, C, D, alpha, beta, X, x_factor, Y): + # All the .load() happen before any .store(). + # Store the accumulator to Y + Y.store(efc_config.accum()) + # Compute the result with alpha, beta scaling and X factor + result = ( + efc_config.relu(efc_config.accum() * alpha + C.load() * beta) + + X.load() * x_factor + ) + D.store(result) + + # Build GEMM object with EFC configuration: + # TODO: generalize acc_dtype and epi_dtype + gemm = DenseGemmAlphaBeta( + acc_dtype, + epi_dtype, + use_2cta_instrs, + mma_tiler_mn, + cluster_shape_mn, + epilogue, + ) + ( + a_tensor, + b_tensor, + a_torch_cpu, + b_torch_cpu, + # The supplemental tensors. + c_tensor, + c_torch_cpu, + c_torch_gpu, + d_tensor, + d_torch_cpu, + d_torch_gpu, + x_tensor, + x_torch_cpu, + x_torch_gpu, + y_tensor, + y_torch_cpu, + y_torch_gpu, + ) = gemm.create_arguments( + l, + m, + n, + k, + a_major, + b_major, + cd_major, + ab_dtype, + # For the supplemental tensors. + c_dtype, + d_dtype, + x_dtype, + y_dtype, + ) + + # Check if the configuration can be implemented. Raise a ValueError + # otherwise. + gemm.check_implementable(a_tensor, b_tensor, d_tensor) + + max_active_clusters = cutlass.utils.HardwareInfo().get_max_active_clusters( + cluster_shape_mn[0] * cluster_shape_mn[1] + ) + + compiled_gemm = gemm.compile( + a_tensor, + b_tensor, + max_active_clusters, + current_stream, + # Here are the supplemental arguments in the same order as for the + # epilogue configuration function. + c_tensor, + d_tensor, + alpha, + beta, + x_tensor, + x_factor, + y_tensor, + # Not really useful here but this is an example of how to pass CuTe + # DSL compilation options. + options="--opt-level=3 --enable-assertions --generate-line-info", + ) + + compiled_gemm( + a_tensor, + b_tensor, + current_stream, + # Here are the supplemental arguments in the same order as for the + # epilogue configuration function. + c_tensor, + d_tensor, + alpha, + beta, + x_tensor, + x_factor, + y_tensor, + ) + + # TODO: unify with modern way to do benchmarking. + exec_time = testing.benchmark( + compiled_gemm, + kernel_arguments=testing.JitArguments( + a_tensor, + b_tensor, + current_stream, + # Here are the supplemental arguments in the same order as for the + # epilogue configuration function. + c_tensor, + d_tensor, + alpha, + beta, + x_tensor, + x_factor, + y_tensor, + ), + stream=current_stream, + warmup_iterations=warmup_iterations, + iterations=iterations, + ) + + print(f"Execution time: {exec_time} us") + + # Compute reference result + if not skip_ref_check: + gemm.compare( + # The usual arguments. + a_torch_cpu, + b_torch_cpu, + epi_dtype, + tolerance, + # For the tensor check. + c_torch_gpu, + d_torch_gpu, + x_torch_gpu, + y_torch_gpu, + # The EFC epilogue arguments. + c_torch_cpu, + d_torch_cpu, + alpha, + beta, + x_torch_cpu, + x_factor, + y_torch_cpu, + ) + + +if __name__ == "__main__": + args = DenseGemmAlphaBeta.CLIParser().parse() + + try: + run( + args.mnkl, + args.ab_dtype, + args.acc_dtype, + args.epi_dtype, + args.a_major, + args.b_major, + args.cd_major, + args.c_dtype, + args.d_dtype, + args.alpha, + args.beta, + args.x_dtype, + args.x_factor, + args.y_dtype, + args.mma_tiler_mn, + args.cluster_shape_mn, + args.use_2cta_instrs, + args.tolerance, + args.warmup_iterations, + args.iterations, + args.skip_ref_check, + ) + print("PASS") + except Exception as exc: + traceback.print_exception(exc) + raise diff --git a/examples/python/CuTeDSL/blackwell/epilogue/synthetic_custom_epilogue_dense_gemm.py b/examples/python/CuTeDSL/blackwell/epilogue/synthetic_custom_epilogue_dense_gemm.py new file mode 100644 index 00000000..0516cebb --- /dev/null +++ b/examples/python/CuTeDSL/blackwell/epilogue/synthetic_custom_epilogue_dense_gemm.py @@ -0,0 +1,403 @@ +# 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. + +import traceback +import typing + +import cuda.bindings.driver as cuda +import torch + +import cutlass +import cutlass.cute.testing as testing +import cutlass.torch as cutlass_torch + +from common_dense_gemm_efc import DenseGemmEFC +import common_efc + +""" +A high-performance persistent batched dense GEMM (D = alpha * A * B + beta * C) example for the NVIDIA Blackwell SM100 architecture +using CUTE DSL. +- Matrix A is MxKxL, L is batch dimension, A can be row-major("K") or column-major("M") +- Matrix B is NxKxL, L is batch dimension, B can be row-major("N") or column-major("K") +- Matrix C is MxNxL, L is batch dimension, C can be row-major("N") or column-major("M") +- Matrix D is MxNxL, L is batch dimension, D can be row-major("N") or column-major("M") +- alpha and beta are float scalars + +This GEMM kernel supports the following features: + - Utilizes Tensor Memory Access (TMA) for efficient memory operations + - Utilizes Blackwell's tcgen05.mma for matrix multiply-accumulate (MMA) operations (including 2cta mma instructions) + - Implements TMA multicast with cluster to reduce L2 memory traffic + - Support persistent tile scheduling to better overlap memory load/store with mma between tiles + - Support warp specialization to avoid explicit pipelining between mainloop load and mma + +This GEMM works as follows: +1. DMA warp: Load A and B matrices from global memory (GMEM) to shared memory (SMEM) using TMA operations. +2. MMA warp: Perform matrix multiply-accumulate (MMA) operations using tcgen05.mma instruction. +3. EPILOGUE warp: + - Load completed accumulator from tensor memory (TMEM) to registers (RMEM) using tcgen05.ld. + - Load C matrix from global memory (GMEM) to shared memory (SMEM) using TMA operations and then copied to registers (RMEM). + - Compute D = alpha * accumulator + beta * C. + - Type convert D matrix to output type. + - Store D matrix from registers (RMEM) to shared memory (SMEM) to global memory (GMEM) with TMA operations, + - Optionally accept an elementwise lambda function epilogue_op to apply to the output tensor: + e.g., relu can set epilogue_op = lambda x: cute.where(x > 0, x, cute.full_like(x, 0)) + +SM100 tcgen05.mma instructions operate as follows: +- Read matrix A from SMEM +- Read matrix B from SMEM +- Write accumulator to TMEM +The accumulator in TMEM must then be loaded to registers before writing back to GMEM. + +Input arguments to this example is same as dense_gemm.py. + +.. code-block:: bash + + python examples/internal/blackwell/epilogue/synthetic_custom_epilogue_dense_gemm.py \ + --ab_dtype Float16 --c_dtype Float16 --d_dtype Float16 --acc_dtype Float32 --epi_dtype Float32 \ + --mma_tiler_mn 256,128 --cluster_shape_mn 2,1 \ + --mnkl 8192,8192,8192,1 \ + --use_2cta_instrs --alpha 2.0 --beta 1.0 --t_dtype Float32 --read_tensors 2 --written_tensors 3 + +To collect performance with NCU profiler: + +.. code-block:: bash + + ncu python examples/internal/blackwell/epilogue/synthetic_custom_epilogue_dense_gemm.py \ + --ab_dtype Float16 --c_dtype Float16 --d_dtype Float16 --acc_dtype Float32 --epi_dtype Float32 \ + --mma_tiler_mn 256,128 --cluster_shape_mn 2,1 \ + --mnkl 8192,8192,8192,1 \ + --use_2cta_instrs --alpha 2.0 --beta 1.0 --t_dtype Float32 --read_tensors 2 --written_tensors 3 \ + --warmup_iterations 1 --iterations 10 --skip_ref_check + + +Constraints are same as dense_gemm.py: +* Supported input data types: fp16, bf16, tf32, int8, uint8, fp8 (e4m3fn, e5m2), + see detailed valid dtype combinations in below SM100PersistentDenseGemmAlphaBetaKernel class documentation +* A/B tensor must have the same data type +* C/D tensor must have the same major order +* Mma tiler M must be 64/128 (use_2cta_instrs=False) or 128/256 (use_2cta_instrs=True) +* Mma tiler N must be 32-256, step 32 +* Cluster shape M/N must be positive and power of 2, total cluster size <= 16 +* Cluster shape M must be multiple of 2 if use_2cta_instrs=True +* The contiguous dimension of A/B/C/D tensors must be at least 16 bytes aligned, + i.e, number of elements is a multiple of 4, 8, and 16 for TFloat32, + Float16/BFloat16, and Int8/Uint8/Float8, respectively. +* OOB tiles are not allowed when TMA store is disabled +""" + + +def format_as_cli_args( + mnkl: typing.Tuple[int, int, int, int], + ab_dtype: typing.Type[cutlass.Numeric], + acc_dtype: typing.Type[cutlass.Numeric], + epi_dtype: typing.Type[cutlass.Numeric], + a_major: str, + b_major: str, + cd_major: str, + mma_tiler_mn: typing.Tuple[int, int], + cluster_shape_mn: typing.Tuple[int, int], + use_2cta_instrs: bool, + t_dtype: typing.Type[cutlass.Numeric], + alpha: float, + beta: float, + read_tensors: int, + written_tensors: int, + tolerance: float, +) -> str: + """Format test parameters as CLI arguments for synthetic_custom_epilogue_dense_gemm.py""" + + # Get base command from DenseGemmEFC class + base_cmd = DenseGemmEFC.format_as_cli_args( + "synthetic_custom_epilogue_dense_gemm.py", + mnkl, + ab_dtype, + acc_dtype, + epi_dtype, + a_major, + b_major, + cd_major, + mma_tiler_mn, + cluster_shape_mn, + use_2cta_instrs, + tolerance, + ) + + # Add synthetic epilogue-specific parameters + specific_args = ( + f" --alpha {alpha} " + f"--beta {beta} " + f"--t_dtype {DenseGemmEFC.dtype_name(t_dtype)} " + f"--read_tensors {read_tensors} " + f"--written_tensors {written_tensors}" + ) + + return base_cmd + specific_args + + +def run( + mnkl: typing.Tuple[int, int, int, int], + ab_dtype: typing.Type[cutlass.Numeric], + acc_dtype: typing.Type[cutlass.Numeric], + epi_dtype: typing.Type[cutlass.Numeric], + a_major: str, + b_major: str, + cd_major: str, + alpha: float, + beta: float, + t_dtype: typing.Type[cutlass.Numeric], + mma_tiler_mn: typing.Tuple[int, int], + cluster_shape_mn: typing.Tuple[int, int], + use_2cta_instrs: bool, + tolerance: float, + warmup_iterations: int = 0, + iterations: int = 1, + skip_ref_check: bool = False, + read_tensors: int = 1, + written_tensors: int = 1, + verbose: bool = False, +): + """ + Prepare A/B/C/D tensors, launch GPU kernel, and reference checking. + """ + print("Running Blackwell Persistent Dense GEMM test with:") + print(f"mnkl: {mnkl}") + print(f"AB dtype: {ab_dtype}, Acc dtype: {acc_dtype}, Epi dtype: {epi_dtype}") + print( + f"Matrix majors - A: {a_major}, B: {b_major}, loaded: {cd_major}, stored: {cd_major}" + ) + print(f"Mma Tiler (M, N): {mma_tiler_mn}, Cluster Shape (M, N): {cluster_shape_mn}") + print(f"2CTA MMA instructions: {'True' if use_2cta_instrs else 'False'}") + print(f"Tolerance: {tolerance}") + print(f"Warmup iterations: {warmup_iterations}") + print(f"Iterations: {iterations}") + print(f"Skip reference checking: {skip_ref_check}") + print("Epilogue:") + print(f"\t{alpha = }, {beta = }") + print(f"\t{t_dtype = !s}") + print(f"\t{read_tensors = }, {written_tensors = }") + + # Unpack parameters + m, n, k, l = mnkl + + if not torch.cuda.is_available(): + raise RuntimeError("A GPU is required to run this example!") + + # Get current CUDA stream from PyTorch + torch_stream = torch.cuda.current_stream() + # Get the raw stream pointer as a CUstream + current_stream = cuda.CUstream(torch_stream.cuda_stream) + + def meta_epilogue(read_tensors, written_tensors): + """Build a synthetic epilogue function with parameters + (self, alpha, beta, read_t0, read_t1,..., read_t{read__tensors-1}, written_t0, written_t1,..., written_t{written_tensors-1}""" + + param_names = ( + ["efc_config", "alpha", "beta"] + + [f"read_t{i}" for i in range(read_tensors)] + + [f"written_t{i}" for i in range(written_tensors)] + ) + + assert written_tensors > 0, ( + "At least one tensor must be written in the epilogue." + ) + + def computation_impl(efc_config, alpha, beta, *tensors): + """Implementation of the epilogue computation.""" + read = beta + for tensor in tensors[:read_tensors]: + read += tensor.load() * alpha + if read_tensors > 0: + # Can use some CuTe/PyTorch-like functions exposed under + # efc_config namespace for portability: + read = efc_config.where( + read < 1, read, read * efc_config.full_like(read, 2) + ) + + t = efc_config.accum() + for tensor in tensors[read_tensors:]: + t = t * alpha + read + 5000 + tensor.store(t) + + # Wrap the implementation with a function with the correct parameter + # names. + return common_efc.create_named_epilogue(param_names, computation_impl) + + epilogue = meta_epilogue(read_tensors, written_tensors) + + # Build GEMM object with EFC configuration: + # TODO: generalize acc_dtype and epi_dtype + gemm = DenseGemmEFC( + acc_dtype, + epi_dtype, + use_2cta_instrs, + mma_tiler_mn, + cluster_shape_mn, + epilogue, + ) + ( + a_tensor, + b_tensor, + a_torch_cpu, + b_torch_cpu, + ) = gemm.create_arguments(l, m, n, k, a_major, b_major, cd_major, ab_dtype) + + # Create all the supplemental tensors. + t_torch_cpu, t_torch_gpu, t_tensor = ([], [], []) + for i in range(read_tensors + written_tensors): + t_torch_cpu.append(cutlass_torch.matrix(l, m, n, cd_major == "m", t_dtype)) + tensor, torch_gpu = cutlass_torch.cute_tensor_like( + t_torch_cpu[i], t_dtype, is_dynamic_layout=True, assumed_align=16 + ) + t_tensor.append(tensor) + t_torch_gpu.append(torch_gpu) + + # Check if configuration can be implemented + gemm.check_implementable(a_tensor, b_tensor, t_tensor[0]) + + max_active_clusters = cutlass.utils.HardwareInfo().get_max_active_clusters( + cluster_shape_mn[0] * cluster_shape_mn[1] + ) + + compiled_gemm = gemm.compile( + a_tensor, + b_tensor, + max_active_clusters, + current_stream, + # Here are the supplemental arguments in the same order as for the + # epilogue configuration function. + alpha, + beta, + *t_tensor, + ) + torch.cuda.synchronize() + + # TODO: unify with modern way to do benchmarking. + exec_time = testing.benchmark( + compiled_gemm, + kernel_arguments=testing.JitArguments( + a_tensor, + b_tensor, + current_stream, + # Here are the supplemental arguments in the same order as for the + # epilogue configuration function. + alpha, + beta, + *t_tensor, + ), + stream=current_stream, + warmup_iterations=warmup_iterations, + iterations=iterations, + ) + + print(f"Execution time: {exec_time} us") + + # Evaluate the epilogue on the host: + gemm.evaluate_on_cpu( + a_torch_cpu, + b_torch_cpu, + epi_dtype, + # The EFC arguments: + alpha, + beta, + *t_torch_cpu, + ) + + # Print tensors if verbose mode is enabled + if verbose: + print("\n=== Read Tensors ===") + for i in range(read_tensors): + print(f"\nRead Tensor {i} (GPU):") + print(t_torch_gpu[i].cpu()) + + print("\n=== Written Tensors ===") + for i in range(written_tensors): + idx = read_tensors + i + print(f"\nWritten Tensor {i} (GPU):") + print(t_torch_gpu[idx].cpu()) + print() + + # Assert close results between the values computed on GPU and CPU. + for torch_gpu, torch_cpu in zip(t_torch_gpu, t_torch_cpu): + torch.testing.assert_close( + torch_gpu.cpu(), torch_cpu, atol=tolerance, rtol=1e-03 + ) + + +if __name__ == "__main__": + cli = DenseGemmEFC.CLIParser() + cli.parser.add_argument( + "--alpha", type=float, default=1.0, help="alpha scale factor" + ) + cli.parser.add_argument("--beta", type=float, default=0.0, help="beta scale factor") + cli.parser.add_argument("--t_dtype", type=cutlass.dtype, default=cutlass.Float32) + cli.parser.add_argument( + "--read_tensors", + type=int, + default=1, + help="number of tensors to read inside the epilogue", + ) + cli.parser.add_argument( + "--written_tensors", + type=int, + default=1, + help="number of tensors to write inside the epilogue", + ) + cli.parser.add_argument( + "--verbose", + action="store_true", + help="print read and written tensors", + ) + args = cli.parse() + + try: + run( + args.mnkl, + args.ab_dtype, + args.acc_dtype, + args.epi_dtype, + args.a_major, + args.b_major, + args.cd_major, + args.alpha, + args.beta, + args.t_dtype, + args.mma_tiler_mn, + args.cluster_shape_mn, + args.use_2cta_instrs, + args.tolerance, + args.warmup_iterations, + args.iterations, + args.skip_ref_check, + args.read_tensors, + args.written_tensors, + args.verbose, + ) + print("PASS") + except Exception as exc: + traceback.print_exception(exc) + raise diff --git a/examples/python/CuTeDSL/blackwell/fmha.py b/examples/python/CuTeDSL/blackwell/fmha.py index f52d9538..7b91051a 100644 --- a/examples/python/CuTeDSL/blackwell/fmha.py +++ b/examples/python/CuTeDSL/blackwell/fmha.py @@ -33,16 +33,14 @@ import sys import time from typing import Type, Tuple, Union, Optional -import torch -import torch.nn.functional as F import cuda.bindings.driver as cuda +import torch import cutlass import cutlass.cute as cute import cutlass.cute.nvgpu.tcgen05 as tcgen05 import cutlass.utils as utils import cutlass.pipeline as pipeline -import cutlass.torch as cutlass_torch import cutlass.utils.blackwell_helpers as sm100_utils import cutlass.cute.testing as testing from cutlass.cute.runtime import from_dlpack @@ -174,8 +172,7 @@ class BlackwellFusedMultiHeadAttentionForward: self.load_warp_id = 13 self.epilogue_warp_id = 14 self.empty_warp_id = 15 - SM100_TMEM_CAPACITY_COLUMNS = 512 - self.tmem_alloc_cols = SM100_TMEM_CAPACITY_COLUMNS + self.tmem_alloc_cols = cute.arch.get_max_tmem_alloc_cols("sm_100") self.threads_per_warp = 32 self.threads_per_cta = self.threads_per_warp * len( @@ -1637,6 +1634,8 @@ class BlackwellFusedMultiHeadAttentionForward: :type atom_args: tuple :param tensor_args: Tuple containing softmax related tensors :type tensor_args: tuple + :param fused_mask: Compute trip counts and apply masking for attention blocks + :type fused_mask: fmha_utils.FusedMask :return: Updated state values (row_max, row_sum, and pipeline related arguments) :rtype: tuple """ @@ -1729,20 +1728,16 @@ class BlackwellFusedMultiHeadAttentionForward: tTMEM_STORErS_x4_e, cute.make_layout(frg_tile) ) for j in range(frg_cnt): - for k in range(0, cute.size(tTMEM_LOADrS_frg, mode=[0]), 2): - tTMEM_LOADrS_frg[k, j], tTMEM_LOADrS_frg[k + 1, j] = ( - cute.arch.fma_packed_f32x2( - (tTMEM_LOADrS_frg[k, j], tTMEM_LOADrS_frg[k + 1, j]), - (scale, scale), - (minus_row_max_scale, minus_row_max_scale), - ) + for k in cutlass.range( + cute.size(tTMEM_LOADrS_frg, mode=[0]), vectorize=True + ): + tTMEM_LOADrS_frg[k, j] = ( + tTMEM_LOADrS_frg[k, j] * scale + minus_row_max_scale ) tTMEM_LOADrS_frg[k, j] = cute.math.exp2( tTMEM_LOADrS_frg[k, j], fastmath=True ) - tTMEM_LOADrS_frg[k + 1, j] = cute.math.exp2( - tTMEM_LOADrS_frg[k + 1, j], fastmath=True - ) + s_vec = tTMEM_LOADrS_frg[None, j].load() tTMEM_STORErS_x4_e_frg[None, j].store(s_vec.to(self.q_dtype)) # Sequence barrier arrive @@ -1859,6 +1854,8 @@ class BlackwellFusedMultiHeadAttentionForward: :type s0_s1_sequence_pipeline: pipeline.PipelineAsync :param tile_sched_params: Parameters for tile scheduling :type tile_sched_params: fmha_utils.FmhaStaticTileSchedulerParams + :param fused_mask: Compute trip counts and apply masking for attention blocks + :type fused_mask: fmha_utils.FusedMask """ tidx, _, _ = cute.arch.thread_idx() thread_idx = tidx % ( @@ -2204,11 +2201,8 @@ class BlackwellFusedMultiHeadAttentionForward: ) cute.copy(tiled_tmem_load, tTMEM_LOADtO_i, tTMrO_i) - for j in range(0, cute.size(tTMrO_i), 2): - tTMrO_i[j], tTMrO_i[j + 1] = cute.arch.mul_packed_f32x2( - (tTMrO_i[j], tTMrO_i[j + 1]), - (scale, scale), - ) + for j in cutlass.range(cute.size(tTMrO_i), vectorize=True): + tTMrO_i[j] = tTMrO_i[j] * scale cute.copy(tiled_tmem_store, tTMrO_i, tTMEM_STOREtO_i) @cute.jit @@ -2310,11 +2304,8 @@ class BlackwellFusedMultiHeadAttentionForward: tTMEM_LOADoO[None, 0, 0, i].shape, self.pv_acc_dtype ) cute.copy(tiled_tmem_load, tTMEM_LOADtO_i, tTMrO) - for j in range(0, cute.size(tTMrO), 2): - tTMrO[j], tTMrO[j + 1] = cute.arch.mul_packed_f32x2( - (tTMrO[j], tTMrO[j + 1]), - (scale, scale), - ) + for j in range(cute.size(tTMrO), vectorize=True): + tTMrO[j] = tTMrO[j] * scale tSMrO = cute.make_rmem_tensor(tTMrO.shape, self.o_dtype) o_vec = tTMrO.load() tSMrO.store(o_vec.to(self.o_dtype)) @@ -2327,7 +2318,10 @@ class BlackwellFusedMultiHeadAttentionForward: mLSE[row_idx + cuseqlen_q, blk_coord[2]] = lse # fence view async shared - cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) def run( @@ -2442,6 +2436,7 @@ def run( print(f" iterations: {iterations}") print(f" skip_ref_check: {skip_ref_check}") print(f" use_cold_l2: {use_cold_l2}") + import cutlass.torch as cutlass_torch # Unpack parameters b, s_q, h_q, d = q_shape @@ -2950,7 +2945,7 @@ def run( else: lse_tensor = None - return testing.JitArguments( + args = testing.JitArguments( q_tensor_workspace.iterator, k_tensor_workspace.iterator, v_tensor_workspace.iterator, @@ -2970,6 +2965,15 @@ def run( ), current_stream, ) + args.add_to_scope( + [ + q_tensor_workspace, + k_tensor_workspace, + v_tensor_workspace, + o_tensor_workspace, + ] + ) + return args workspace_count = 1 if use_cold_l2: diff --git a/examples/python/CuTeDSL/blackwell/fmha_bwd.py b/examples/python/CuTeDSL/blackwell/fmha_bwd.py index f40af955..5725d5e8 100644 --- a/examples/python/CuTeDSL/blackwell/fmha_bwd.py +++ b/examples/python/CuTeDSL/blackwell/fmha_bwd.py @@ -166,8 +166,7 @@ class BlackwellFusedMultiHeadAttentionBackward: self.num_reduce_warps = 4 self.num_compute_warps = 8 - SM100_TMEM_CAPACITY_COLUMNS = 512 - self.tmem_alloc_cols = SM100_TMEM_CAPACITY_COLUMNS + self.tmem_alloc_cols = cute.arch.get_max_tmem_alloc_cols("sm_100") self.threads_per_warp = 32 self.threads_per_cta = self.threads_per_warp * ( @@ -2093,7 +2092,6 @@ class BlackwellFusedMultiHeadAttentionBackward: cute.arch.fence_view_async_tmem_load() self.compute_sync_barrier.arrive_and_wait() - cute.arch.fence_view_async_tmem_load() cute.copy(tiled_r2t, tRT_rST_reshaped, tRT_tP) @@ -2161,7 +2159,10 @@ class BlackwellFusedMultiHeadAttentionBackward: cute.autovec_copy(tTR_rdST, sdS_slice) # Notify for dS - cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) compute_mma_dS_pipeline.producer_commit(compute_mma_dS_producer_state) compute_mma_dS_producer_state.advance() @@ -2279,7 +2280,10 @@ class BlackwellFusedMultiHeadAttentionBackward: ) # Wait for the stores to all be visible to the TMA - cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) self.reduce_sync_barrier.arrive_and_wait() if warp_idx == 0: @@ -2489,11 +2493,13 @@ class BlackwellFusedMultiHeadAttentionBackward: mma_compute_dKdV_pipeline.consumer_wait(mma_compute_dKdV_consumer_state) + # Load tdKtdK cute.copy(tiled_t2r_dK, tTR_tdK, tTR_rdK) for i in cutlass.range(cute.size(tTR_rdK), unroll_full=True): tTR_rdK[i] = scale_softmax * tTR_rdK[i] + # Store tdKgdK self.store(tTR_gdK, tTR_rdK, tTR_cdK, (K, D)) cute.arch.fence_view_async_tmem_load() diff --git a/examples/python/CuTeDSL/blackwell/grouped_blockscaled_gemm.py b/examples/python/CuTeDSL/blackwell/grouped_blockscaled_gemm.py index 7cc59105..0f2dd5fa 100644 --- a/examples/python/CuTeDSL/blackwell/grouped_blockscaled_gemm.py +++ b/examples/python/CuTeDSL/blackwell/grouped_blockscaled_gemm.py @@ -193,8 +193,7 @@ class Sm100GroupedBlockScaledGemmKernel: num_threads=64, ) self.smem_capacity = utils.get_smem_capacity_in_bytes("sm_100") - SM100_TMEM_CAPACITY_COLUMNS = 512 - self.num_tmem_alloc_cols = SM100_TMEM_CAPACITY_COLUMNS + self.num_tmem_alloc_cols = cute.arch.get_max_tmem_alloc_cols("sm_100") # Set up configurations that dependent on gemm inputs. def _setup_attributes(self): @@ -423,6 +422,7 @@ class Sm100GroupedBlockScaledGemmKernel: self.b_dtype = initial_b.element_type self.sf_dtype = initial_sfa.element_type self.c_dtype = initial_c.element_type + self.is_nvfp4_output = self.c_dtype is cutlass.Float4E2M1FN self.a_major_mode = utils.LayoutEnum.from_tensor(initial_a).mma_major_mode() self.b_major_mode = utils.LayoutEnum.from_tensor(initial_b).mma_major_mode() self.c_layout = utils.LayoutEnum.from_tensor(initial_c) @@ -893,7 +893,7 @@ class Sm100GroupedBlockScaledGemmKernel: cute.group_modes(tCgB, 0, 3), ) - # TMA Load SFA partition_S/D + # TMA load scaled factor A partition_S/D sfa_cta_layout = a_cta_layout # ((atom_v, rest_v), STAGE) # ((atom_v, rest_v), RestM, RestK, RestL) @@ -907,7 +907,7 @@ class Sm100GroupedBlockScaledGemmKernel: tAsSFA = cute.filter_zeros(tAsSFA) tAgSFA = cute.filter_zeros(tAgSFA) - # TMA Load SFB partition_S/D + # TMA load scaled factor B partition_S/D sfb_cta_layout = cute.make_layout( cute.slice_(cluster_layout_sfb_vmnk, (0, None, 0, 0)).shape ) @@ -970,222 +970,239 @@ class Sm100GroupedBlockScaledGemmKernel: tensormaps[(tensormap_workspace_idx, 4, None)].iterator ) + # + # Persistent tile scheduling loop + # + # When the problem shapes are on device, we launch one CTA per SM. + # The if condition later prevents the warps from extra CTAs from doing any work. + tile_sched = utils.StaticPersistentGroupTileScheduler.create( + tile_sched_params, + cute.arch.block_idx(), + grid_dim, + self.cluster_tile_shape_mnk, + utils.create_initial_search_state(), + group_count, + problem_sizes_mnkl, + ) + initial_work_tile_info = tile_sched.initial_work_tile_info() + # # Specialized TMA load warp # - if warp_idx == self.tma_warp_id: + if warp_idx == self.tma_warp_id and initial_work_tile_info.is_valid_tile: # # Persistent tile scheduling loop # - tile_sched = utils.StaticPersistentTileScheduler.create( - tile_sched_params, cute.arch.block_idx(), grid_dim - ) - # grouped gemm tile scheduler helper will compute the group index for the tile we're working on - group_gemm_ts_helper = utils.GroupedGemmTileSchedulerHelper( - group_count, - tile_sched_params, - self.cluster_tile_shape_mnk, - utils.create_initial_search_state(), - ) + work_tile = initial_work_tile_info + tensormap_init_done = cutlass.Boolean(False) # group index of last tile last_group_idx = cutlass.Int32(-1) - work_tile = tile_sched.initial_work_tile_info() - ab_producer_state = pipeline.make_pipeline_state( pipeline.PipelineUserType.Producer, self.num_ab_stage ) while work_tile.is_valid_tile: - cur_tile_coord = work_tile.tile_idx - grouped_gemm_cta_tile_info = group_gemm_ts_helper.delinearize_z( - cur_tile_coord, - problem_sizes_mnkl, - ) + grouped_gemm_cta_tile_info = work_tile.group_search_result cur_k_tile_cnt = grouped_gemm_cta_tile_info.cta_tile_count_k cur_group_idx = grouped_gemm_cta_tile_info.group_idx - is_group_changed = cur_group_idx != last_group_idx - # skip tensormap update if we're working on the same group - if is_group_changed: - real_tensor_a = self.make_tensor_abc_for_tensormap_update( - cur_group_idx, - self.a_dtype, - ( - grouped_gemm_cta_tile_info.problem_shape_m, - grouped_gemm_cta_tile_info.problem_shape_n, - grouped_gemm_cta_tile_info.problem_shape_k, - ), - strides_abc, - ptrs_abc, - 0, # 0 for tensor A - ) - real_tensor_b = self.make_tensor_abc_for_tensormap_update( - cur_group_idx, - self.b_dtype, - ( - grouped_gemm_cta_tile_info.problem_shape_m, - grouped_gemm_cta_tile_info.problem_shape_n, - grouped_gemm_cta_tile_info.problem_shape_k, - ), - strides_abc, - ptrs_abc, - 1, # 1 for tensor B - ) - real_tensor_sfa = self.make_tensor_sfasfb_for_tensormap_update( - cur_group_idx, - self.sf_dtype, - ( - grouped_gemm_cta_tile_info.problem_shape_m, - grouped_gemm_cta_tile_info.problem_shape_n, - grouped_gemm_cta_tile_info.problem_shape_k, - ), - ptrs_sfasfb, - 0, # 0 for tensor SFA - ) - real_tensor_sfb = self.make_tensor_sfasfb_for_tensormap_update( - cur_group_idx, - self.sf_dtype, - ( - grouped_gemm_cta_tile_info.problem_shape_m, - grouped_gemm_cta_tile_info.problem_shape_n, - grouped_gemm_cta_tile_info.problem_shape_k, - ), - ptrs_sfasfb, - 1, # 1 for tensor SFB - ) - if tensormap_init_done == False: - # wait tensormap initialization complete - self.tensormap_ab_init_barrier.arrive_and_wait() - tensormap_init_done = True + is_k_tile_cnt_zero = cur_k_tile_cnt == 0 + # Do not load any data if cur_k_tile_cnt is 0 + if not is_k_tile_cnt_zero: + is_group_changed = cur_group_idx != last_group_idx + # skip tensormap update if we're working on the same group + if is_group_changed: + real_tensor_a = self.make_tensor_abc_for_tensormap_update( + cur_group_idx, + self.a_dtype, + ( + grouped_gemm_cta_tile_info.problem_shape_m, + grouped_gemm_cta_tile_info.problem_shape_n, + grouped_gemm_cta_tile_info.problem_shape_k, + ), + strides_abc, + ptrs_abc, + 0, # 0 for tensor A + ) + real_tensor_b = self.make_tensor_abc_for_tensormap_update( + cur_group_idx, + self.b_dtype, + ( + grouped_gemm_cta_tile_info.problem_shape_m, + grouped_gemm_cta_tile_info.problem_shape_n, + grouped_gemm_cta_tile_info.problem_shape_k, + ), + strides_abc, + ptrs_abc, + 1, # 1 for tensor B + ) + real_tensor_sfa = self.make_tensor_sfasfb_for_tensormap_update( + cur_group_idx, + self.sf_dtype, + ( + grouped_gemm_cta_tile_info.problem_shape_m, + grouped_gemm_cta_tile_info.problem_shape_n, + grouped_gemm_cta_tile_info.problem_shape_k, + ), + ptrs_sfasfb, + 0, # 0 for tensor SFA + ) + real_tensor_sfb = self.make_tensor_sfasfb_for_tensormap_update( + cur_group_idx, + self.sf_dtype, + ( + grouped_gemm_cta_tile_info.problem_shape_m, + grouped_gemm_cta_tile_info.problem_shape_n, + grouped_gemm_cta_tile_info.problem_shape_k, + ), + ptrs_sfasfb, + 1, # 1 for tensor SFB + ) + if not tensormap_init_done: + # wait tensormap initialization complete + self.tensormap_ab_init_barrier.arrive_and_wait() + tensormap_init_done = True - tensormap_manager.update_tensormap( - ( - real_tensor_a, - real_tensor_b, - real_tensor_sfa, - real_tensor_sfb, - ), - (tma_atom_a, tma_atom_b, tma_atom_sfa, tma_atom_sfb), - ( - tensormap_a_gmem_ptr, - tensormap_b_gmem_ptr, - tensormap_sfa_gmem_ptr, - tensormap_sfb_gmem_ptr, - ), - self.tma_warp_id, - ( - tensormap_a_smem_ptr, - tensormap_b_smem_ptr, - tensormap_sfa_smem_ptr, - tensormap_sfb_smem_ptr, - ), + tensormap_manager.update_tensormap( + ( + real_tensor_a, + real_tensor_b, + real_tensor_sfa, + real_tensor_sfb, + ), + (tma_atom_a, tma_atom_b, tma_atom_sfa, tma_atom_sfb), + ( + tensormap_a_gmem_ptr, + tensormap_b_gmem_ptr, + tensormap_sfa_gmem_ptr, + tensormap_sfb_gmem_ptr, + ), + self.tma_warp_id, + ( + tensormap_a_smem_ptr, + tensormap_b_smem_ptr, + tensormap_sfa_smem_ptr, + tensormap_sfb_smem_ptr, + ), + ) + + mma_tile_coord_mnl = ( + grouped_gemm_cta_tile_info.cta_tile_idx_m + // cute.size(tiled_mma.thr_id.shape), + grouped_gemm_cta_tile_info.cta_tile_idx_n, + 0, ) - mma_tile_coord_mnl = ( - grouped_gemm_cta_tile_info.cta_tile_idx_m - // cute.size(tiled_mma.thr_id.shape), - grouped_gemm_cta_tile_info.cta_tile_idx_n, - 0, - ) + # + # Slice to per mma tile index + # + # ((atom_v, rest_v), RestK) + tAgA_slice = tAgA[ + (None, mma_tile_coord_mnl[0], None, mma_tile_coord_mnl[2]) + ] + # ((atom_v, rest_v), RestK) + tBgB_slice = tBgB[ + (None, mma_tile_coord_mnl[1], None, mma_tile_coord_mnl[2]) + ] - # - # Slice to per mma tile index - # - # ((atom_v, rest_v), RestK) - tAgA_slice = tAgA[ - (None, mma_tile_coord_mnl[0], None, mma_tile_coord_mnl[2]) - ] - # ((atom_v, rest_v), RestK) - tBgB_slice = tBgB[ - (None, mma_tile_coord_mnl[1], None, mma_tile_coord_mnl[2]) - ] + # ((atom_v, rest_v), RestK) + tAgSFA_slice = tAgSFA[ + (None, mma_tile_coord_mnl[0], None, mma_tile_coord_mnl[2]) + ] + # ((atom_v, rest_v), RestK) + tBgSFB_slice = tBgSFB[ + (None, mma_tile_coord_mnl[1], None, mma_tile_coord_mnl[2]) + ] - # ((atom_v, rest_v), RestK) - tAgSFA_slice = tAgSFA[ - (None, mma_tile_coord_mnl[0], None, mma_tile_coord_mnl[2]) - ] - # ((atom_v, rest_v), RestK) - tBgSFB_slice = tBgSFB[ - (None, mma_tile_coord_mnl[1], None, mma_tile_coord_mnl[2]) - ] - - # Peek (try_wait) AB buffer empty for k_tile = prefetch_k_tile_cnt - ab_producer_state.reset_count() - peek_ab_empty_status = cutlass.Boolean(1) - if ab_producer_state.count < cur_k_tile_cnt: - peek_ab_empty_status = ab_pipeline.producer_try_acquire( - ab_producer_state - ) - - if is_group_changed: - tensormap_manager.fence_tensormap_update(tensormap_a_gmem_ptr) - tensormap_manager.fence_tensormap_update(tensormap_b_gmem_ptr) - tensormap_manager.fence_tensormap_update(tensormap_sfa_gmem_ptr) - tensormap_manager.fence_tensormap_update(tensormap_sfb_gmem_ptr) - # - # Tma load loop - # - for k_tile in cutlass.range(0, cur_k_tile_cnt, 1, unroll=1): - # Conditionally wait for AB buffer empty - ab_pipeline.producer_acquire( - ab_producer_state, peek_ab_empty_status - ) - - # TMA load A/B/SFA/SFB - cute.copy( - tma_atom_a, - tAgA_slice[(None, ab_producer_state.count)], - tAsA[(None, ab_producer_state.index)], - tma_bar_ptr=ab_pipeline.producer_get_barrier(ab_producer_state), - mcast_mask=a_full_mcast_mask, - tma_desc_ptr=tensormap_manager.get_tensormap_ptr( - tensormap_a_gmem_ptr, - cute.AddressSpace.generic, - ), - ) - cute.copy( - tma_atom_b, - tBgB_slice[(None, ab_producer_state.count)], - tBsB[(None, ab_producer_state.index)], - tma_bar_ptr=ab_pipeline.producer_get_barrier(ab_producer_state), - mcast_mask=b_full_mcast_mask, - tma_desc_ptr=tensormap_manager.get_tensormap_ptr( - tensormap_b_gmem_ptr, - cute.AddressSpace.generic, - ), - ) - cute.copy( - tma_atom_sfa, - tAgSFA_slice[(None, ab_producer_state.count)], - tAsSFA[(None, ab_producer_state.index)], - tma_bar_ptr=ab_pipeline.producer_get_barrier(ab_producer_state), - mcast_mask=sfa_full_mcast_mask, - tma_desc_ptr=tensormap_manager.get_tensormap_ptr( - tensormap_sfa_gmem_ptr, - cute.AddressSpace.generic, - ), - ) - cute.copy( - tma_atom_sfb, - tBgSFB_slice[(None, ab_producer_state.count)], - tBsSFB[(None, ab_producer_state.index)], - tma_bar_ptr=ab_pipeline.producer_get_barrier(ab_producer_state), - mcast_mask=sfb_full_mcast_mask, - tma_desc_ptr=tensormap_manager.get_tensormap_ptr( - tensormap_sfb_gmem_ptr, - cute.AddressSpace.generic, - ), - ) - - # Peek (try_wait) AB buffer empty for k_tile = prefetch_k_tile_cnt + k_tile + 1 - ab_producer_state.advance() + # Peek (try_wait) AB buffer empty for k_tile = prefetch_k_tile_cnt + ab_producer_state.reset_count() peek_ab_empty_status = cutlass.Boolean(1) if ab_producer_state.count < cur_k_tile_cnt: peek_ab_empty_status = ab_pipeline.producer_try_acquire( ab_producer_state ) + if is_group_changed: + tensormap_manager.fence_tensormap_update(tensormap_a_gmem_ptr) + tensormap_manager.fence_tensormap_update(tensormap_b_gmem_ptr) + tensormap_manager.fence_tensormap_update(tensormap_sfa_gmem_ptr) + tensormap_manager.fence_tensormap_update(tensormap_sfb_gmem_ptr) + # + # Tma load loop + # + for k_tile in cutlass.range(0, cur_k_tile_cnt, 1, unroll=1): + # Conditionally wait for AB buffer empty + ab_pipeline.producer_acquire( + ab_producer_state, peek_ab_empty_status + ) + + # TMA load A/B/SFA/SFB + cute.copy( + tma_atom_a, + tAgA_slice[(None, ab_producer_state.count)], + tAsA[(None, ab_producer_state.index)], + tma_bar_ptr=ab_pipeline.producer_get_barrier( + ab_producer_state + ), + mcast_mask=a_full_mcast_mask, + tma_desc_ptr=tensormap_manager.get_tensormap_ptr( + tensormap_a_gmem_ptr, + cute.AddressSpace.generic, + ), + ) + cute.copy( + tma_atom_b, + tBgB_slice[(None, ab_producer_state.count)], + tBsB[(None, ab_producer_state.index)], + tma_bar_ptr=ab_pipeline.producer_get_barrier( + ab_producer_state + ), + mcast_mask=b_full_mcast_mask, + tma_desc_ptr=tensormap_manager.get_tensormap_ptr( + tensormap_b_gmem_ptr, + cute.AddressSpace.generic, + ), + ) + cute.copy( + tma_atom_sfa, + tAgSFA_slice[(None, ab_producer_state.count)], + tAsSFA[(None, ab_producer_state.index)], + tma_bar_ptr=ab_pipeline.producer_get_barrier( + ab_producer_state + ), + mcast_mask=sfa_full_mcast_mask, + tma_desc_ptr=tensormap_manager.get_tensormap_ptr( + tensormap_sfa_gmem_ptr, + cute.AddressSpace.generic, + ), + ) + cute.copy( + tma_atom_sfb, + tBgSFB_slice[(None, ab_producer_state.count)], + tBsSFB[(None, ab_producer_state.index)], + tma_bar_ptr=ab_pipeline.producer_get_barrier( + ab_producer_state + ), + mcast_mask=sfb_full_mcast_mask, + tma_desc_ptr=tensormap_manager.get_tensormap_ptr( + tensormap_sfb_gmem_ptr, + cute.AddressSpace.generic, + ), + ) + + # Peek (try_wait) AB buffer empty for k_tile = prefetch_k_tile_cnt + k_tile + 1 + ab_producer_state.advance() + peek_ab_empty_status = cutlass.Boolean(1) + if ab_producer_state.count < cur_k_tile_cnt: + peek_ab_empty_status = ab_pipeline.producer_try_acquire( + ab_producer_state + ) + else: + if not tensormap_init_done: + # wait tensormap initialization complete + self.tensormap_ab_init_barrier.arrive_and_wait() + tensormap_init_done = True # # Advance to next tile # @@ -1201,7 +1218,7 @@ class Sm100GroupedBlockScaledGemmKernel: # # Specialized MMA warp # - if warp_idx == self.mma_warp_id: + if warp_idx == self.mma_warp_id and initial_work_tile_info.is_valid_tile: # # Initialize tensormaps for A, B, SFA and SFB # @@ -1279,18 +1296,8 @@ class Sm100GroupedBlockScaledGemmKernel: # # Persistent tile scheduling loop # - tile_sched = utils.StaticPersistentTileScheduler.create( - tile_sched_params, cute.arch.block_idx(), grid_dim - ) - # grouped gemm tile scheduler helper will compute the group index for the tile we're working on - group_gemm_ts_helper = utils.GroupedGemmTileSchedulerHelper( - group_count, - tile_sched_params, - self.cluster_tile_shape_mnk, - utils.create_initial_search_state(), - ) + work_tile = initial_work_tile_info - work_tile = tile_sched.initial_work_tile_info() ab_consumer_state = pipeline.make_pipeline_state( pipeline.PipelineUserType.Consumer, self.num_ab_stage ) @@ -1298,15 +1305,14 @@ class Sm100GroupedBlockScaledGemmKernel: pipeline.PipelineUserType.Producer, self.num_acc_stage ) while work_tile.is_valid_tile: - cur_tile_coord = work_tile.tile_idx + cur_group_idx = work_tile.group_search_result.group_idx + problem_shape_k = work_tile.group_search_result.problem_shape_k + # MMA warp is only interested in number of tiles along K dimension - ( - cur_k_tile_cnt, - cur_group_idx, - ) = group_gemm_ts_helper.search_cluster_tile_count_k( - cur_tile_coord, - problem_sizes_mnkl, - ) + cur_k_tile_cnt = ( + problem_shape_k + self.cluster_tile_shape_mnk[2] - 1 + ) // self.cluster_tile_shape_mnk[2] + is_k_tile_cnt_zero = cur_k_tile_cnt == 0 # (MMA, MMA_M, MMA_N) tCtAcc = tCtAcc_base[(None, None, None, acc_producer_state.index)] @@ -1322,7 +1328,7 @@ class Sm100GroupedBlockScaledGemmKernel: # # Wait for accumulator buffer empty # - if is_leader_cta: + if is_leader_cta and not is_k_tile_cnt_zero: acc_pipeline.producer_acquire(acc_producer_state) # @@ -1408,9 +1414,10 @@ class Sm100GroupedBlockScaledGemmKernel: # # Async arrive accumulator buffer full # - if is_leader_cta: - acc_pipeline.producer_commit(acc_producer_state) - acc_producer_state.advance() + if not is_k_tile_cnt_zero: + if is_leader_cta: + acc_pipeline.producer_commit(acc_producer_state) + acc_producer_state.advance() # # Advance to next tile @@ -1426,7 +1433,7 @@ class Sm100GroupedBlockScaledGemmKernel: # # Specialized epilogue warps # - if warp_idx < self.mma_warp_id: + if warp_idx < self.mma_warp_id and initial_work_tile_info.is_valid_tile: # initialize tensorap for C tensormap_manager.init_tensormap_from_atom( tma_atom_c, @@ -1483,18 +1490,7 @@ class Sm100GroupedBlockScaledGemmKernel: # # Persistent tile scheduling loop # - tile_sched = utils.StaticPersistentTileScheduler.create( - tile_sched_params, cute.arch.block_idx(), grid_dim - ) - # grouped gemm tile scheduler helper will compute the group index for the tile we're working on - group_gemm_ts_helper = utils.GroupedGemmTileSchedulerHelper( - group_count, - tile_sched_params, - self.cluster_tile_shape_mnk, - utils.create_initial_search_state(), - ) - - work_tile = tile_sched.initial_work_tile_info() + work_tile = initial_work_tile_info acc_consumer_state = pipeline.make_pipeline_state( pipeline.PipelineUserType.Consumer, self.num_acc_stage @@ -1513,14 +1509,13 @@ class Sm100GroupedBlockScaledGemmKernel: last_group_idx = cutlass.Int32(-1) while work_tile.is_valid_tile: - cur_tile_coord = work_tile.tile_idx - grouped_gemm_cta_tile_info = group_gemm_ts_helper.delinearize_z( - cur_tile_coord, - problem_sizes_mnkl, - ) + grouped_gemm_cta_tile_info = work_tile.group_search_result cur_group_idx = grouped_gemm_cta_tile_info.group_idx + cur_k_tile_cnt = grouped_gemm_cta_tile_info.cta_tile_count_k + is_k_tile_cnt_zero = cur_k_tile_cnt == 0 is_group_changed = cur_group_idx != last_group_idx + # We still need to store 0s when k_tile_cnt is 0 if is_group_changed: # construct tensor c based on real shape, stride information real_tensor_c = self.make_tensor_abc_for_tensormap_update( @@ -1549,7 +1544,6 @@ class Sm100GroupedBlockScaledGemmKernel: grouped_gemm_cta_tile_info.cta_tile_idx_n, 0, ) - cur_k_tile_cnt = grouped_gemm_cta_tile_info.cta_tile_count_k # # Slice to per mma tile index @@ -1573,7 +1567,8 @@ class Sm100GroupedBlockScaledGemmKernel: # # Wait for accumulator buffer full # - acc_pipeline.consumer_wait(acc_consumer_state) + if not is_k_tile_cnt_zero: + acc_pipeline.consumer_wait(acc_consumer_state) tTR_tAcc = cute.group_modes(tTR_tAcc, 3, cute.rank(tTR_tAcc)) bSG_gC = cute.group_modes(bSG_gC, 1, cute.rank(bSG_gC)) @@ -1588,17 +1583,34 @@ class Sm100GroupedBlockScaledGemmKernel: subtile_cnt = cute.size(tTR_tAcc.shape, mode=[3]) num_prev_subtiles = tile_sched.num_tiles_executed * subtile_cnt for subtile_idx in range(subtile_cnt): - # - # Load accumulator from tensor memory buffer to register - # - tTR_tAcc_mn = tTR_tAcc[(None, None, None, subtile_idx)] - cute.copy(tiled_copy_t2r, tTR_tAcc_mn, tTR_rAcc) + if not is_k_tile_cnt_zero: + # + # Load accumulator from tensor memory buffer to register + # + tTR_tAcc_mn = tTR_tAcc[(None, None, None, subtile_idx)] + cute.copy(tiled_copy_t2r, tTR_tAcc_mn, tTR_rAcc) - # - # Convert to C type - # - acc_vec = tiled_copy_r2s.retile(tTR_rAcc).load() - tRS_rC.store(acc_vec.to(self.c_dtype)) + # + # Convert to C type + # + acc_vec = tiled_copy_r2s.retile(tTR_rAcc).load() + tRS_rC.store(acc_vec.to(self.c_dtype)) + else: + if cutlass.const_expr(self.is_nvfp4_output): + zeros_i8 = cute.make_rmem_tensor( + cute.recast_layout( + cutlass.Int8.width, + self.c_dtype.width, + tRS_rC.layout, + ), + cutlass.Int8, + ) + zeros_i8.fill(0) + tRS_rC.store( + cute.recast_tensor(zeros_i8, self.c_dtype).load() + ) + else: + tRS_rC.fill(0) # # Store C to shared memory @@ -1610,7 +1622,10 @@ class Sm100GroupedBlockScaledGemmKernel: tRS_sC[(None, None, None, c_buffer)], ) # Fence and barrier to make sure shared memory store is visible to TMA store - cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) self.epilog_sync_barrier.arrive_and_wait() # @@ -1633,9 +1648,10 @@ class Sm100GroupedBlockScaledGemmKernel: # # Async arrive accumulator buffer empty # - with cute.arch.elect_one(): - acc_pipeline.consumer_release(acc_consumer_state) - acc_consumer_state.advance() + if not is_k_tile_cnt_zero: + with cute.arch.elect_one(): + acc_pipeline.consumer_release(acc_consumer_state) + acc_consumer_state.advance() # # Advance to next tile @@ -2442,13 +2458,6 @@ def create_tensor_and_stride( torch_tensor_cpu, dtype, is_dynamic_layout, assumed_align=16 ) - # Mark tensor with element divisibility for 16B alignment - cute_tensor.mark_compact_shape_dynamic( - mode=0 if is_mode0_major else 1, - stride_order=(2, 1, 0) if is_mode0_major else (2, 0, 1), - divisibility=32 if dtype == cutlass.Float4E2M1FN else 16, - ) - # omit stride for L mode as it is always 1 stride = (1, mode0) if is_mode0_major else (mode1, 1) @@ -2552,7 +2561,7 @@ def create_scale_factor_tensor(l, mn, k, sf_vec_size, dtype): def ceil_div(a, b): return (a + b - 1) // b - sf_k = ceil_div(k, sf_vec_size) + sf_k = max(1, ceil_div(k, sf_vec_size)) ref_shape = (l, mn, sf_k) atom_m = (32, 4) @@ -2675,6 +2684,7 @@ def create_tensors_sfasfb_for_all_groups( def run( num_groups: int, problem_sizes_mnkl: List[Tuple[int, int, int, int]], + host_problem_shape_available: bool, ab_dtype: Type[cutlass.Numeric], sf_dtype: Type[cutlass.Numeric], sf_vec_size: int, @@ -2761,22 +2771,40 @@ def run( sf_vec_size, ) - # Choose A, B, C, SFA, SFB with the smallest size to create initial tensormaps - key_size_a = lambda item: item[1][0] * item[1][2] - key_size_b = lambda item: item[1][1] * item[1][2] - key_size_c = lambda item: item[1][0] * item[1][1] - # Find the indices of the groups with the smallest tensor sizes - min_a_idx, _ = min(enumerate(problem_sizes_mnkl), key=key_size_a) - min_b_idx, _ = min(enumerate(problem_sizes_mnkl), key=key_size_b) - min_c_idx, _ = min(enumerate(problem_sizes_mnkl), key=key_size_c) + # Setup inital tensors for TMA of A,B and C + alignment = 16 # 16 bytes aligned + divisibility_ab = 32 if ab_dtype == cutlass.Float4E2M1FN else 16 + divisibility_c = 32 if c_dtype == cutlass.Float4E2M1FN else 16 + divisibility_sf = 32 if sf_dtype == cutlass.Float4E2M1FN else 16 + + min_ab_size = alignment * 8 // ab_dtype.width # alignment bytes of width + div_mul_ab = (divisibility_ab + min_ab_size - 1) // min_ab_size + min_ab_size = min_ab_size * div_mul_ab + + min_c_size = alignment * 8 // c_dtype.width + div_mul_c = (divisibility_c + min_c_size - 1) // min_c_size + min_c_size = min_c_size * div_mul_c + + min_sf_size = alignment * 8 // sf_dtype.width + div_mul_sf = (divisibility_sf + min_sf_size - 1) // min_sf_size + min_sf_size = min_sf_size * div_mul_sf + initial_cute_tensors_abc = [ - cute_tensors_abc[min_a_idx][0], # A with smallest (m, k) - cute_tensors_abc[min_b_idx][1], # B with smallest (n, k) - cute_tensors_abc[min_c_idx][2], # C with smallest (m, n) + create_tensor_and_stride(1, min_ab_size, min_ab_size, a_major == "m", ab_dtype)[ + 2 + ], + create_tensor_and_stride(1, min_ab_size, min_ab_size, b_major == "n", ab_dtype)[ + 2 + ], + create_tensor_and_stride(1, min_c_size, min_c_size, c_major == "m", c_dtype)[2], ] initial_cute_tensors_sfasfb = [ - cute_tensors_sfasfb[min_a_idx][0], # SFA with smallest (m, k)'s group - cute_tensors_sfasfb[min_b_idx][1], # SFB with smallest (n, k)'s group + create_tensor_and_stride(1, min_sf_size, min_sf_size, a_major == "m", sf_dtype)[ + 2 + ], + create_tensor_and_stride(1, min_sf_size, min_sf_size, b_major == "n", sf_dtype)[ + 2 + ], ] hardware_info = cutlass.utils.HardwareInfo() @@ -2867,6 +2895,19 @@ def run( # Initialize Stream current_stream = cutlass_torch.default_stream() + # If the host problem shape is available, we will launch the grid with only + # the necessary clusters. The function compute_total_num_clusters() does that. + # If the problem shape only exists on device, we will need to launch all active + # clusters possible on a device. + if host_problem_shape_available: + print("Problem shapes available on host and device") + total_num_clusters = compute_total_num_clusters( + problem_sizes_mnkl, cluster_tile_shape_mn + ) + else: + print("Problem shapes available only on device") + total_num_clusters = max_active_clusters + # Compile grouped GEMM kernel compiled_grouped_gemm = cute.compile( grouped_blockscaled_gemm, @@ -2980,18 +3021,23 @@ def run( ) initial_cute_tensors_abc_workspace = [ - cute_tensors_abc_workspace[min_a_idx][0], # A with smallest (m, k) - cute_tensors_abc_workspace[min_b_idx][1], # B with smallest (n, k) - cute_tensors_abc_workspace[min_c_idx][2], # C with smallest (m, n) + create_tensor_and_stride( + 1, min_ab_size, min_ab_size, a_major == "m", ab_dtype + )[2], + create_tensor_and_stride( + 1, min_ab_size, min_ab_size, b_major == "n", ab_dtype + )[2], + create_tensor_and_stride( + 1, min_c_size, min_c_size, c_major == "m", c_dtype + )[2], ] - initial_cute_tensors_sfasfb_workspace = [ - cute_tensors_sfasfb_workspace[min_a_idx][ - 0 - ], # SFA with smallest (m, k)'s group - cute_tensors_sfasfb_workspace[min_b_idx][ - 1 - ], # SFB with smallest (n, k)'s group + create_tensor_and_stride( + 1, min_sf_size, min_sf_size, a_major == "m", sf_dtype + )[2], + create_tensor_and_stride( + 1, min_sf_size, min_sf_size, b_major == "n", sf_dtype + )[2], ] # Create new tensors for this workspace @@ -3022,7 +3068,7 @@ def run( is_dynamic_layout=False, ) - return cute.testing.JitArguments( + args = cute.testing.JitArguments( initial_cute_tensors_abc_workspace[0], initial_cute_tensors_abc_workspace[1], initial_cute_tensors_abc_workspace[2], @@ -3035,6 +3081,8 @@ def run( tensormap_workspace, current_stream, ) + args.add_to_scope([torch_tensors_abc_workspace, torch_tensors_sfasfb_workspace]) + return args workspace_count = 1 if use_cold_l2: @@ -3078,6 +3126,18 @@ def run( iterations=iterations, ) + runtime_s = exec_time / 1.0e6 + fmas = 0 + for group in range(num_groups): + [M, N, K, _] = problem_sizes_mnkl[group] + fmas += M * N * K + flop = 2 * fmas + gflop = flop / 1.0e9 + gflops = gflop / runtime_s + + print("Average Runtime : ", exec_time / 1000, "ms") + print("GFLOPS : ", gflops) + return exec_time # Return execution time in microseconds @@ -3138,6 +3198,11 @@ if __name__ == "__main__": default=(128, 128), help="Mma tile shape (comma-separated)", ) + parser.add_argument( + "--host_problem_shape_available", + action="store_true", + help="Enable the compute of grid based upon host problem shape", + ) parser.add_argument( "--cluster_shape_mn", type=parse_comma_separated_ints, @@ -3195,6 +3260,7 @@ if __name__ == "__main__": run( args.num_groups, args.problem_sizes_mnkl, + args.host_problem_shape_available, args.ab_dtype, args.sf_dtype, args.sf_vec_size, diff --git a/examples/python/CuTeDSL/blackwell/grouped_gemm.py b/examples/python/CuTeDSL/blackwell/grouped_gemm.py index 1c63ac8b..50c208f2 100644 --- a/examples/python/CuTeDSL/blackwell/grouped_gemm.py +++ b/examples/python/CuTeDSL/blackwell/grouped_gemm.py @@ -335,9 +335,11 @@ class GroupedGemmKernel: :type stream: cuda.CUstream :raises TypeError: If A and B data types do not match. """ + self.a_dtype = initial_a.element_type self.b_dtype = initial_b.element_type self.c_dtype = initial_c.element_type + self.a_major_mode = utils.LayoutEnum.from_tensor(initial_a).mma_major_mode() self.b_major_mode = utils.LayoutEnum.from_tensor(initial_b).mma_major_mode() self.c_layout = utils.LayoutEnum.from_tensor(initial_c) @@ -475,6 +477,7 @@ class GroupedGemmKernel: block=[self.threads_per_cta, 1, 1], cluster=(*self.cluster_shape_mn, 1), stream=stream, + min_blocks_per_mp=1, ) return @@ -553,28 +556,38 @@ class GroupedGemmKernel: tensormap_c_smem_ptr = ( tensormap_b_smem_ptr + GroupedGemmKernel.bytes_per_tensormap // 8 ) - ab_full_mbar_ptr = storage.ab_full_mbar_ptr.data_ptr() - ab_empty_mbar_ptr = storage.ab_empty_mbar_ptr.data_ptr() - acc_full_mbar_ptr = storage.acc_full_mbar_ptr.data_ptr() - acc_empty_mbar_ptr = storage.acc_empty_mbar_ptr.data_ptr() # init barrier for loading A, B with TMA - if warp_idx == self.epilog_warp_id[0]: - for k_stage in range(self.num_ab_stage): - num_tma_producer = self.num_mcast_ctas_a + self.num_mcast_ctas_b - 1 - with cute.arch.elect_one(): - cute.arch.mbarrier_init(ab_full_mbar_ptr + k_stage, 1) - cute.arch.mbarrier_init( - ab_empty_mbar_ptr + k_stage, num_tma_producer - ) + ab_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) + num_tma_producer = self.num_mcast_ctas_a + self.num_mcast_ctas_b - 1 + ab_pipeline_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, num_tma_producer + ) + ab_pipeline = pipeline.PipelineTmaUmma.create( + barrier_storage=storage.ab_full_mbar_ptr.data_ptr(), + num_stages=self.num_ab_stage, + producer_group=ab_pipeline_producer_group, + consumer_group=ab_pipeline_consumer_group, + tx_count=self.num_tma_load_bytes, + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ) # Accumulator barrier init - if warp_idx == self.mma_warp_id: - for acc_stage in range(self.num_acc_stage): - with cute.arch.elect_one(): - cute.arch.mbarrier_init(acc_full_mbar_ptr + acc_stage, 1) - cute.arch.mbarrier_init( - acc_empty_mbar_ptr + acc_stage, 8 if use_2cta_instrs else 4 - ) + acc_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) + num_acc_consumer_threads = len(self.epilog_warp_id) * ( + 2 if use_2cta_instrs else 1 + ) + acc_pipeline_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, num_acc_consumer_threads + ) + acc_pipeline = pipeline.PipelineUmmaAsync.create( + barrier_storage=storage.acc_full_mbar_ptr.data_ptr(), + num_stages=self.num_acc_stage, + producer_group=acc_pipeline_producer_group, + consumer_group=acc_pipeline_consumer_group, + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ) # Tensor memory dealloc barrier init tmem = utils.TmemAllocator( storage.tmem_holding_buf, @@ -747,10 +760,26 @@ class GroupedGemmKernel: tensormap_b_init_ptr = tensormap_b_ptr tensormap_c_init_ptr = tensormap_c_ptr + # + # Persistent tile scheduling loop + # + # When the problem shapes are on device, we launch one CTA per SM. + # The if condition later prevents the warps from extra CTAs from doing any work. + tile_sched = utils.StaticPersistentGroupTileScheduler.create( + tile_sched_params, + bid, + grid_dim, + self.cluster_tile_shape_mnk, + utils.create_initial_search_state(), + group_count, + problem_sizes_mnkl, + ) + initial_work_tile_info = tile_sched.initial_work_tile_info() + # # Specialized TMA load warp # - if warp_idx == self.tma_warp_id: + if warp_idx == self.tma_warp_id and initial_work_tile_info.is_valid_tile: # Initialize tensormaps for A, B if cutlass.const_expr(self.delegate_tensormap_ab_init == False): tensormap_manager.init_tensormap_from_atom( @@ -759,185 +788,161 @@ class GroupedGemmKernel: tensormap_manager.init_tensormap_from_atom( tma_atom_b, tensormap_b_init_ptr, self.tma_warp_id ) - # - # Persistent tile scheduling loop - # - tile_sched = utils.StaticPersistentTileScheduler.create( - tile_sched_params, bid, grid_dim - ) - # grouped gemm tile scheduler helper will compute the group index for the tile we're working on - group_gemm_ts_helper = utils.GroupedGemmTileSchedulerHelper( - group_count, - tile_sched_params, - self.cluster_tile_shape_mnk, - utils.create_initial_search_state(), - ) + tensormap_init_done = cutlass.Boolean(False) - # tile count we have searched - total_k_tile_cnt = cutlass.Int32(0) # group index of last tile last_group_idx = cutlass.Int32(-1) - work_tile = tile_sched.initial_work_tile_info() + + work_tile = initial_work_tile_info + ab_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.num_ab_stage + ) + while work_tile.is_valid_tile: - cur_tile_coord = work_tile.tile_idx - grouped_gemm_cta_tile_info = group_gemm_ts_helper.delinearize_z( - cur_tile_coord, - problem_sizes_mnkl, - ) + grouped_gemm_cta_tile_info = work_tile.group_search_result + cur_k_tile_cnt = grouped_gemm_cta_tile_info.cta_tile_count_k + is_k_tile_cnt_zero = cur_k_tile_cnt == 0 cur_group_idx = grouped_gemm_cta_tile_info.group_idx - is_group_changed = cur_group_idx != last_group_idx - # skip tensormap update if we're working on the same group - if is_group_changed: - real_tensor_a = self.make_tensor_for_tensormap_update( - cur_group_idx, - self.a_dtype, - ( - grouped_gemm_cta_tile_info.problem_shape_m, - grouped_gemm_cta_tile_info.problem_shape_n, - grouped_gemm_cta_tile_info.problem_shape_k, - ), - strides_abc, - ptrs_abc, - 0, # 0 for tensor A + # Do not load any data if cur_k_tile_cnt is 0 + if not is_k_tile_cnt_zero: + is_group_changed = cur_group_idx != last_group_idx + # skip tensormap update if we're working on the same group + if is_group_changed: + real_tensor_a = self.make_tensor_for_tensormap_update( + cur_group_idx, + self.a_dtype, + ( + grouped_gemm_cta_tile_info.problem_shape_m, + grouped_gemm_cta_tile_info.problem_shape_n, + grouped_gemm_cta_tile_info.problem_shape_k, + ), + strides_abc, + ptrs_abc, + 0, # 0 for tensor A + ) + real_tensor_b = self.make_tensor_for_tensormap_update( + cur_group_idx, + self.b_dtype, + ( + grouped_gemm_cta_tile_info.problem_shape_m, + grouped_gemm_cta_tile_info.problem_shape_n, + grouped_gemm_cta_tile_info.problem_shape_k, + ), + strides_abc, + ptrs_abc, + 1, # 1 for tensor B + ) + # wait tensormap initialization complete before update + if not tensormap_init_done: + if cutlass.const_expr(self.delegate_tensormap_ab_init): + self.tensormap_ab_init_barrier.arrive_and_wait() + tensormap_manager.fence_tensormap_initialization() + tensormap_init_done = True + + tensormap_manager.update_tensormap( + (real_tensor_a, real_tensor_b), + (tma_atom_a, tma_atom_b), + (tensormap_a_ptr, tensormap_b_ptr), + self.tma_warp_id, + (tensormap_a_smem_ptr, tensormap_b_smem_ptr), + ) + + mma_tile_coord_mnl = ( + grouped_gemm_cta_tile_info.cta_tile_idx_m + // cute.size(tiled_mma.thr_id.shape), + grouped_gemm_cta_tile_info.cta_tile_idx_n, + 0, ) - real_tensor_b = self.make_tensor_for_tensormap_update( - cur_group_idx, - self.b_dtype, - ( - grouped_gemm_cta_tile_info.problem_shape_m, - grouped_gemm_cta_tile_info.problem_shape_n, - grouped_gemm_cta_tile_info.problem_shape_k, - ), - strides_abc, - ptrs_abc, - 1, # 1 for tensor B - ) - # wait tensormap initialization complete before update - if tensormap_init_done == False: + + # + # Slice to per mma tile index + # + # ((atom_v, rest_v), RestK) + tAgA_slice = tAgA[ + (None, mma_tile_coord_mnl[0], None, mma_tile_coord_mnl[2]) + ] + # ((atom_v, rest_v), RestK) + tBgB_slice = tBgB[ + (None, mma_tile_coord_mnl[1], None, mma_tile_coord_mnl[2]) + ] + + # Peek (try_wait) AB buffer empty for k_tile = prefetch_k_tile_cnt + ab_producer_state.reset_count() + peek_ab_empty_status = cutlass.Boolean(1) + if ab_producer_state.count < cur_k_tile_cnt: + peek_ab_empty_status = ab_pipeline.producer_try_acquire( + ab_producer_state + ) + # ensure the update to tensormap has completed before using it + if is_group_changed: + tensormap_manager.fence_tensormap_update(tensormap_a_ptr) + tensormap_manager.fence_tensormap_update(tensormap_b_ptr) + # + # Tma load loop + # + for k_tile in cutlass.range(0, cur_k_tile_cnt, 1, unroll=1): + # Wait for AB buffer empty + ab_pipeline.producer_acquire( + ab_producer_state, peek_ab_empty_status + ) + + # Load A/B with TMA + cute.copy( + tma_atom_a, + tAgA_slice[(None, ab_producer_state.count)], + tAsA[(None, ab_producer_state.index)], + tma_bar_ptr=ab_pipeline.producer_get_barrier( + ab_producer_state + ), + mcast_mask=a_full_mcast_mask, + tma_desc_ptr=tensormap_manager.get_tensormap_ptr( + tensormap_a_ptr, + cute.AddressSpace.generic, + ), + ) + cute.copy( + tma_atom_b, + tBgB_slice[(None, ab_producer_state.count)], + tBsB[(None, ab_producer_state.index)], + tma_bar_ptr=ab_pipeline.producer_get_barrier( + ab_producer_state + ), + mcast_mask=b_full_mcast_mask, + tma_desc_ptr=tensormap_manager.get_tensormap_ptr( + tensormap_b_ptr, + cute.AddressSpace.generic, + ), + ) + + # Peek (try_wait) AB buffer empty for k_tile = prefetch_k_tile_cnt + k_tile + 1 + ab_producer_state.advance() + peek_ab_empty_status = cutlass.Boolean(1) + if ab_producer_state.count < cur_k_tile_cnt: + peek_ab_empty_status = ab_pipeline.producer_try_acquire( + ab_producer_state + ) + else: + # If tensormap initialization is not done, wait for it to complete + if not tensormap_init_done: if cutlass.const_expr(self.delegate_tensormap_ab_init): self.tensormap_ab_init_barrier.arrive_and_wait() tensormap_manager.fence_tensormap_initialization() tensormap_init_done = True - - tensormap_manager.update_tensormap( - (real_tensor_a, real_tensor_b), - (tma_atom_a, tma_atom_b), - (tensormap_a_ptr, tensormap_b_ptr), - self.tma_warp_id, - (tensormap_a_smem_ptr, tensormap_b_smem_ptr), - ) - - mma_tile_coord_mnl = ( - grouped_gemm_cta_tile_info.cta_tile_idx_m - // cute.size(tiled_mma.thr_id.shape), - grouped_gemm_cta_tile_info.cta_tile_idx_n, - 0, - ) - - # - # Slice to per mma tile index - # - # ((atom_v, rest_v), RestK) - tAgA_slice = tAgA[ - (None, mma_tile_coord_mnl[0], None, mma_tile_coord_mnl[2]) - ] - # ((atom_v, rest_v), RestK) - tBgB_slice = tBgB[ - (None, mma_tile_coord_mnl[1], None, mma_tile_coord_mnl[2]) - ] - - num_prev_k_blk = total_k_tile_cnt - total_k_tile_cnt += cur_k_tile_cnt - - # Peek (try_wait) AB buffer empty for k_tile = prefetch_k_tile_cnt - tma_wr_k_tile = cutlass.Int32(0) - smem_wr_buffer = (num_prev_k_blk + tma_wr_k_tile) % self.num_ab_stage - tma_wr_ab_empty_phase = ( - num_prev_k_blk + tma_wr_k_tile - ) // self.num_ab_stage % 2 ^ 1 - peek_ab_empty_status = cute.arch.mbarrier_conditional_try_wait( - tma_wr_k_tile < cur_k_tile_cnt, - ab_empty_mbar_ptr + smem_wr_buffer, - tma_wr_ab_empty_phase, - ) - # ensure the update to tensormap has completed before using it - if is_group_changed: - tensormap_manager.fence_tensormap_update(tensormap_a_ptr) - tensormap_manager.fence_tensormap_update(tensormap_b_ptr) - # - # Tma load loop - # - for k_tile in cutlass.range(0, cur_k_tile_cnt, 1, unroll=1): - tma_wr_k_tile_next = tma_wr_k_tile + 1 - smem_wr_buffer_next = ( - num_prev_k_blk + tma_wr_k_tile_next - ) % self.num_ab_stage - tma_wr_ab_empty_phase_next = ( - tma_wr_ab_empty_phase ^ 1 - if smem_wr_buffer_next == 0 - else tma_wr_ab_empty_phase - ) - - smem_full_mbar_ptr = ab_full_mbar_ptr + smem_wr_buffer - - # Wait for AB buffer empty - if peek_ab_empty_status == 0: - cute.arch.mbarrier_wait( - ab_empty_mbar_ptr + smem_wr_buffer, tma_wr_ab_empty_phase - ) - - # Arrive AB buffer and expect full transaction bytes - if is_leader_cta: - with cute.arch.elect_one(): - cute.arch.mbarrier_arrive_and_expect_tx( - smem_full_mbar_ptr, self.num_tma_load_bytes - ) - - # Load A/B with TMA - cute.copy( - tma_atom_a, - tAgA_slice[(None, tma_wr_k_tile)], - tAsA[(None, smem_wr_buffer)], - tma_bar_ptr=smem_full_mbar_ptr, - mcast_mask=a_full_mcast_mask, - tma_desc_ptr=tensormap_manager.get_tensormap_ptr( - tensormap_a_ptr, - cute.AddressSpace.generic, - ), - ) - cute.copy( - tma_atom_b, - tBgB_slice[(None, tma_wr_k_tile)], - tBsB[(None, smem_wr_buffer)], - tma_bar_ptr=smem_full_mbar_ptr, - mcast_mask=b_full_mcast_mask, - tma_desc_ptr=tensormap_manager.get_tensormap_ptr( - tensormap_b_ptr, - cute.AddressSpace.generic, - ), - ) - - # Peek (try_wait) AB buffer empty for k_tile = prefetch_k_tile_cnt + k_tile + 1 - peek_ab_empty_status = cute.arch.mbarrier_conditional_try_wait( - tma_wr_k_tile_next < cur_k_tile_cnt, - ab_empty_mbar_ptr + smem_wr_buffer_next, - tma_wr_ab_empty_phase_next, - ) - - tma_wr_k_tile = tma_wr_k_tile_next - smem_wr_buffer = smem_wr_buffer_next - tma_wr_ab_empty_phase = tma_wr_ab_empty_phase_next - # Advance to next tile tile_sched.advance_to_next_work() work_tile = tile_sched.get_current_work() last_group_idx = cur_group_idx + # + # Wait A/B buffer empty + # + ab_pipeline.producer_tail(ab_producer_state) + # # Specialized MMA warp # - if warp_idx == self.mma_warp_id: + if warp_idx == self.mma_warp_id and initial_work_tile_info.is_valid_tile: # Bar sync for retrieve tmem ptr from shared mem tmem.wait_for_alloc() @@ -951,63 +956,42 @@ class GroupedGemmKernel: # # Persistent tile scheduling loop # - tile_sched = utils.StaticPersistentTileScheduler.create( - tile_sched_params, bid, grid_dim + work_tile = initial_work_tile_info + ab_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_ab_stage ) - # grouped gemm tile scheduler helper will compute the group index for the tile we're working on - group_gemm_ts_helper = utils.GroupedGemmTileSchedulerHelper( - group_count, - tile_sched_params, - self.cluster_tile_shape_mnk, - utils.create_initial_search_state(), + acc_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.num_acc_stage ) - work_tile = tile_sched.initial_work_tile_info() # tile count we have searched - total_k_tile_cnt = cutlass.Int32(0) while work_tile.is_valid_tile: - cur_tile_coord = work_tile.tile_idx - # MMA warp is only interested in number of tiles along K dimension - ( - cur_k_tile_cnt, - cur_group_idx, - ) = group_gemm_ts_helper.search_cluster_tile_count_k( - cur_tile_coord, - problem_sizes_mnkl, - ) - # Set tensor memory buffer for current tile - acc_buf_idx = tile_sched.num_tiles_executed % self.num_acc_stage - # (MMA, MMA_M, MMA_N) - tCtAcc = tCtAcc_base[(None, None, None, acc_buf_idx)] + cur_group_idx = work_tile.group_search_result.group_idx + problem_shape_k = work_tile.group_search_result.problem_shape_k - num_prev_k_blk = total_k_tile_cnt - total_k_tile_cnt += cur_k_tile_cnt + # MMA warp is only interested in number of tiles along K dimension + cur_k_tile_cnt = ( + problem_shape_k + self.cluster_tile_shape_mnk[2] - 1 + ) // self.cluster_tile_shape_mnk[2] + is_k_tile_cnt_zero = cur_k_tile_cnt == 0 + + # (MMA, MMA_M, MMA_N) + tCtAcc = tCtAcc_base[(None, None, None, acc_producer_state.index)] # Peek (try_wait) AB buffer full for k_tile = 0 - mma_rd_k_tile = cutlass.Int32(0) - smem_rd_buffer = (num_prev_k_blk + mma_rd_k_tile) % self.num_ab_stage + ab_consumer_state.reset_count() + peek_ab_full_status = cutlass.Boolean(1) if is_leader_cta: - need_check_rd_buffer_full = ( - mma_rd_k_tile < cur_k_tile_cnt and is_leader_cta - ) - mma_rd_ab_full_phase = ( - (num_prev_k_blk + mma_rd_k_tile) // self.num_ab_stage % 2 - ) - peek_ab_full_status = cute.arch.mbarrier_conditional_try_wait( - need_check_rd_buffer_full, - ab_full_mbar_ptr + smem_rd_buffer, - mma_rd_ab_full_phase, - ) + if ab_consumer_state.count < cur_k_tile_cnt: + peek_ab_full_status = ab_pipeline.consumer_try_wait( + ab_consumer_state + ) # # Wait for accumulator buffer empty # - acc_empty_phase = ( - tile_sched.num_tiles_executed // self.num_acc_stage % 2 ^ 1 - ) - cute.arch.mbarrier_wait( - acc_empty_mbar_ptr + acc_buf_idx, acc_empty_phase - ) + if not is_k_tile_cnt_zero: + acc_pipeline.producer_acquire(acc_producer_state) # # Reset the ACCUMULATE field for each tile @@ -1017,26 +1001,20 @@ class GroupedGemmKernel: # # Mma mainloop # - for k_tile in range(cur_k_tile_cnt): - mma_rd_k_tile_next = cutlass.Int32(k_tile + 1) - smem_rd_buffer_next = ( - num_prev_k_blk + mma_rd_k_tile_next - ) % self.num_ab_stage - mma_rd_ab_full_phase_next = ( - mma_rd_ab_full_phase ^ 1 - if smem_rd_buffer_next == 0 - else mma_rd_ab_full_phase - ) + for k_tile in cutlass.range(0, cur_k_tile_cnt, 1, unroll=1): # Wait for AB buffer full - if peek_ab_full_status == 0: - cute.arch.mbarrier_wait( - ab_full_mbar_ptr + smem_rd_buffer, mma_rd_ab_full_phase - ) - + ab_pipeline.consumer_wait( + ab_consumer_state, peek_ab_full_status + ) # tCtAcc += tCrA * tCrB num_kblocks = cute.size(tCrA, mode=[2]) for kblock_idx in cutlass.range(num_kblocks, unroll_full=True): - kblock_coord = (None, None, kblock_idx, smem_rd_buffer) + kblock_coord = ( + None, + None, + kblock_idx, + ab_consumer_state.index, + ) cute.gemm( tiled_mma, @@ -1049,48 +1027,37 @@ class GroupedGemmKernel: tiled_mma.set(tcgen05.Field.ACCUMULATE, True) # Async arrive AB buffer empty - with cute.arch.elect_one(): - tcgen05.commit( - ab_empty_mbar_ptr + smem_rd_buffer, - ab_empty_mcast_mask, - self.cta_group, - ) + ab_pipeline.consumer_release(ab_consumer_state) # Peek (try_wait) AB buffer full for k_tile = k_tile + 1 - need_check_rd_buffer_full = ( - mma_rd_k_tile_next < cur_k_tile_cnt and is_leader_cta - ) - - peek_ab_full_status = cute.arch.mbarrier_conditional_try_wait( - need_check_rd_buffer_full, - ab_full_mbar_ptr + smem_rd_buffer_next, - mma_rd_ab_full_phase_next, - ) - - mma_rd_k_tile = mma_rd_k_tile_next - smem_rd_buffer = smem_rd_buffer_next - mma_rd_ab_full_phase = mma_rd_ab_full_phase_next + ab_consumer_state.advance() + peek_ab_full_status = cutlass.Boolean(1) + if ab_consumer_state.count < cur_k_tile_cnt: + peek_ab_full_status = ab_pipeline.consumer_try_wait( + ab_consumer_state + ) # # Async arrive accumulator buffer full # - with cute.arch.elect_one(): - tcgen05.commit( - acc_full_mbar_ptr + acc_buf_idx, - acc_full_mcast_mask, - self.cta_group, - ) + if not is_k_tile_cnt_zero: + acc_pipeline.producer_commit(acc_producer_state) + acc_producer_state.advance() # # Advance to next tile # tile_sched.advance_to_next_work() work_tile = tile_sched.get_current_work() + # + # Wait for accumulator buffer empty + # + acc_pipeline.producer_tail(acc_producer_state) # # Specialized epilogue warps # - if warp_idx < self.mma_warp_id: + if warp_idx < self.mma_warp_id and initial_work_tile_info.is_valid_tile: # initialize tensormap A, B for TMA warp if cutlass.const_expr(self.delegate_tensormap_ab_init): tensormap_manager.init_tensormap_from_atom( @@ -1147,32 +1114,32 @@ class GroupedGemmKernel: # # Persistent tile scheduling loop # - tile_sched = utils.StaticPersistentTileScheduler.create( - tile_sched_params, bid, grid_dim - ) - # grouped gemm tile scheduler helper will compute the group index for the tile we're working on - group_gemm_ts_helper = utils.GroupedGemmTileSchedulerHelper( - group_count, - tile_sched_params, - self.cluster_tile_shape_mnk, - utils.create_initial_search_state(), - ) - work_tile = tile_sched.initial_work_tile_info() + work_tile = initial_work_tile_info + # wait tensormap initialization complete before update tensormap_manager.fence_tensormap_initialization() - # tile count we have searched - total_k_tile_cnt = cutlass.Int32(0) + acc_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_acc_stage + ) + # Threads/warps participating in tma store pipeline + c_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + 32 * len(self.epilog_warp_id), + ) + c_pipeline = pipeline.PipelineTmaStore.create( + num_stages=self.num_epi_stage, + producer_group=c_producer_group, + ) # group index of last tile last_group_idx = cutlass.Int32(-1) while work_tile.is_valid_tile: - cur_tile_coord = work_tile.tile_idx - grouped_gemm_cta_tile_info = group_gemm_ts_helper.delinearize_z( - cur_tile_coord, - problem_sizes_mnkl, - ) + grouped_gemm_cta_tile_info = work_tile.group_search_result cur_group_idx = grouped_gemm_cta_tile_info.group_idx + cur_k_tile_cnt = grouped_gemm_cta_tile_info.cta_tile_count_k + is_k_tile_cnt_zero = cur_k_tile_cnt == 0 is_group_changed = cur_group_idx != last_group_idx + # We still need to store 0s when k_tile_cnt is 0 if is_group_changed: # construct tensor C based on real address, shape and stride information real_tensor_c = self.make_tensor_for_tensormap_update( @@ -1201,8 +1168,6 @@ class GroupedGemmKernel: grouped_gemm_cta_tile_info.cta_tile_idx_n, 0, ) - cur_k_tile_cnt = grouped_gemm_cta_tile_info.cta_tile_count_k - total_k_tile_cnt += cur_k_tile_cnt # # Slice to per mma tile index @@ -1216,17 +1181,16 @@ class GroupedGemmKernel: *mma_tile_coord_mnl, ) ] - - # Set tensor memory buffer for current tile - acc_buf_idx = tile_sched.num_tiles_executed % self.num_acc_stage # (T2R, T2R_M, T2R_N, EPI_M, EPI_M) - tTR_tAcc = tTR_tAcc_base[(None, None, None, None, None, acc_buf_idx)] - + tTR_tAcc = tTR_tAcc_base[ + (None, None, None, None, None, acc_consumer_state.index) + ] # # Wait for accumulator buffer full # - acc_full_phase = tile_sched.num_tiles_executed // self.num_acc_stage % 2 - cute.arch.mbarrier_wait(acc_full_mbar_ptr + acc_buf_idx, acc_full_phase) + # Not waiting for accumulator buffer full when k_tile_cnt is 0 + if not is_k_tile_cnt_zero: + acc_pipeline.consumer_wait(acc_consumer_state) tTR_tAcc = cute.group_modes(tTR_tAcc, 3, cute.rank(tTR_tAcc)) bSG_gC = cute.group_modes(bSG_gC, 1, cute.rank(bSG_gC)) @@ -1240,28 +1204,34 @@ class GroupedGemmKernel: subtile_cnt = cute.size(tTR_tAcc.shape, mode=[3]) num_prev_subtiles = tile_sched.num_tiles_executed * subtile_cnt for subtile_idx in range(subtile_cnt): - # - # Load accumulator from tensor memory buffer to register - # - tTR_tAcc_mn = tTR_tAcc[(None, None, None, subtile_idx)] - cute.copy(tiled_copy_t2r, tTR_tAcc_mn, tTR_rAcc) - - # - # Convert to output type - # - acc_vec = tiled_copy_r2s.retile(tTR_rAcc).load() - tRS_rC.store(acc_vec.to(self.c_dtype)) # # Store C to shared memory # epi_buffer = (num_prev_subtiles + subtile_idx) % self.num_epi_stage + # + # Load accumulator from tensor memory buffer to register + # + tTR_tAcc_mn = tTR_tAcc[(None, None, None, subtile_idx)] + if not is_k_tile_cnt_zero: + cute.copy(tiled_copy_t2r, tTR_tAcc_mn, tTR_rAcc) + + # + # Convert to output type + # + acc_vec = tiled_copy_r2s.retile(tTR_rAcc).load() + tRS_rC.store(acc_vec.to(self.c_dtype)) + else: + tRS_rC.fill(0) cute.copy( tiled_copy_r2s, tRS_rC, tRS_sC[(None, None, None, epi_buffer)], ) # Fence and barrier to make sure shared memory store is visible to TMA store - cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) self.epilog_sync_barrier.arrive_and_wait() # # store C to global memory with TMA @@ -1276,19 +1246,17 @@ class GroupedGemmKernel: cute.AddressSpace.generic, ), ) - cute.arch.cp_async_bulk_commit_group() - cute.arch.cp_async_bulk_wait_group( - self.num_epi_stage - 1, read=True - ) + # Fence and barrier to make sure shared memory store is visible to TMA store + c_pipeline.producer_commit() + c_pipeline.producer_acquire() self.epilog_sync_barrier.arrive_and_wait() # # Async arrive accumulator buffer empty # - with cute.arch.elect_one(): - cute.arch.mbarrier_arrive( - acc_empty_mbar_ptr + acc_buf_idx, - cta_rank_in_cluster // 2 * 2 if use_2cta_instrs else None, - ) + if not is_k_tile_cnt_zero: + with cute.arch.elect_one(): + acc_pipeline.consumer_release(acc_consumer_state) + acc_consumer_state.advance() # # Advance to next tile @@ -1305,13 +1273,9 @@ class GroupedGemmKernel: tmem.free(tmem_ptr) # - # Wait a/b buffer empty + # Wait for C store complete # - if warp_idx == self.epilog_warp_id[0]: - cute.arch.mbarrier_wait( - (ab_empty_mbar_ptr + ((total_k_tile_cnt - 1) % self.num_ab_stage)), - (((total_k_tile_cnt - 1) // self.num_ab_stage) % 2), - ) + c_pipeline.producer_tail() @cute.jit def make_tensor_for_tensormap_update( @@ -1649,7 +1613,7 @@ class GroupedGemmKernel: problem_shape_ntile_mnl, (*cluster_shape_mn, 1) ) - grid = utils.StaticPersistentTileScheduler.get_grid_shape( + grid = utils.StaticPersistentGroupTileScheduler.get_grid_shape( tile_sched_params, max_active_clusters ) @@ -1866,6 +1830,7 @@ def create_tensors_for_all_groups( def run( num_groups: int, problem_sizes_mnkl: tuple[int, int, int, int], + host_problem_shape_available: bool, ab_dtype: Type[cutlass.Numeric], c_dtype: Type[cutlass.Numeric], acc_dtype: Type[cutlass.Numeric], @@ -1975,18 +1940,18 @@ def run( c_major, ) - # Choose A, B, C with the smallest size to create initial tensormaps - key_size_a = lambda item: item[1][0] * item[1][2] - key_size_b = lambda item: item[1][1] * item[1][2] - key_size_c = lambda item: item[1][0] * item[1][1] - # Find the indices of the groups with the smallest tensor sizes - min_a_idx, _ = min(enumerate(problem_sizes_mnkl), key=key_size_a) - min_b_idx, _ = min(enumerate(problem_sizes_mnkl), key=key_size_b) - min_c_idx, _ = min(enumerate(problem_sizes_mnkl), key=key_size_c) + # Setup inital tensors for TMA of A,B and C + alignment = 16 # 16 bytes aligned + min_ab_size = alignment * 8 // ab_dtype.width + min_c_size = alignment * 8 // c_dtype.width initial_cute_tensors_abc = [ - cute_tensors_abc[min_a_idx][0], # A with smallest (m, k) - cute_tensors_abc[min_b_idx][1], # B with smallest (n, k) - cute_tensors_abc[min_c_idx][2], # C with smallest (m, n) + create_tensor_and_stride(1, min_ab_size, min_ab_size, a_major == "m", ab_dtype)[ + 2 + ], + create_tensor_and_stride(1, min_ab_size, min_ab_size, b_major == "n", ab_dtype)[ + 2 + ], + create_tensor_and_stride(1, min_c_size, min_c_size, c_major == "m", c_dtype)[2], ] hardware_info = utils.HardwareInfo() @@ -1994,6 +1959,7 @@ def run( max_active_clusters = hardware_info.get_max_active_clusters( cluster_shape_mn[0] * cluster_shape_mn[1] ) + # Prepare tensormap buffer for each SM num_tensormap_buffers = sm_count tensormap_shape = ( @@ -2069,9 +2035,19 @@ def run( cluster_tile_shape_mn = compute_cluster_tile_shape( mma_tiler_mn, cluster_shape_mn, use_2cta_instrs ) - total_num_clusters = compute_total_num_clusters( - problem_sizes_mnkl, cluster_tile_shape_mn - ) + + # If the host problem shape is available, we will launch the grid with only + # the necessary clusters. The function compute_total_num_clusters() does that. + # If the problem shape only exists on device, we will need to launch all active + # clusters possible on a device. + if host_problem_shape_available: + print("Problem shapes available on host and device") + total_num_clusters = compute_total_num_clusters( + problem_sizes_mnkl, cluster_tile_shape_mn + ) + else: + print("Problem shapes available only on device") + total_num_clusters = max_active_clusters # Initialize Stream current_stream = cutlass_torch.default_stream() @@ -2079,6 +2055,7 @@ def run( # try to check CUDA version to decide the opt level try: from cutlass import CUDA_VERSION + opt_level = ( 3 if CUDA_VERSION.major < 13 @@ -2131,6 +2108,9 @@ def run( rtol=1e-05, ) + if iterations <= 0: + return 0 + def generate_tensors(): # Reuse existing CPU tensors and create new GPU tensors from them ( @@ -2150,9 +2130,15 @@ def run( ) initial_cute_tensors_abc_workspace = [ - cute_tensors_abc_workspace[min_a_idx][0], # A with smallest (m, k) - cute_tensors_abc_workspace[min_b_idx][1], # B with smallest (n, k) - cute_tensors_abc_workspace[min_c_idx][2], # C with smallest (m, n) + create_tensor_and_stride( + 1, min_ab_size, min_ab_size, a_major == "m", ab_dtype + )[2], + create_tensor_and_stride( + 1, min_ab_size, min_ab_size, b_major == "n", ab_dtype + )[2], + create_tensor_and_stride( + 1, min_c_size, min_c_size, c_major == "m", c_dtype + )[2], ] # Create new tensors for this workspace @@ -2176,7 +2162,7 @@ def run( is_dynamic_layout=False, ) - return testing.JitArguments( + args = testing.JitArguments( initial_cute_tensors_abc_workspace[0], initial_cute_tensors_abc_workspace[1], initial_cute_tensors_abc_workspace[2], @@ -2186,6 +2172,8 @@ def run( tensormap_workspace, current_stream, ) + args.add_to_scope([torch_tensors_abc_workspace]) + return args workspace_count = 1 if use_cold_l2: @@ -2225,6 +2213,18 @@ def run( iterations=iterations, ) + runtime_s = exec_time / 1.0e6 + fmas = 0 + for group in range(num_groups): + [M, N, K, _] = problem_sizes_mnkl[group] + fmas += M * N * K + flop = 2 * fmas + gflop = flop / 1.0e9 + gflops = gflop / runtime_s + + print("Average Runtime : ", exec_time / 1000, "ms") + print("GFLOPS : ", gflops) + return exec_time # Return execution time in microseconds @@ -2270,15 +2270,20 @@ if __name__ == "__main__": parser.add_argument( "--num_groups", type=int, - default=2, + default=3, help="Number of groups", ) parser.add_argument( "--problem_sizes_mnkl", type=parse_comma_separated_tuples, - default=((128, 128, 128, 1), (128, 128, 128, 1)), + default=((128, 128, 128, 1), (512, 128, 128, 1), (128, 256, 128, 1)), help="a tuple of problem sizes for each group (comma-separated tuples)", ) + parser.add_argument( + "--host_problem_shape_available", + action="store_true", + help="Enable the compute of grid based upon host problem shape", + ) parser.add_argument( "--mma_tiler_mn", type=parse_comma_separated_ints, @@ -2362,6 +2367,7 @@ if __name__ == "__main__": run( args.num_groups, args.problem_sizes_mnkl, + args.host_problem_shape_available, args.ab_dtype, args.c_dtype, args.acc_dtype, diff --git a/examples/python/CuTeDSL/blackwell/mamba2_ssd/mamba2_ssd.py b/examples/python/CuTeDSL/blackwell/mamba2_ssd/mamba2_ssd.py index c72c24c8..630de111 100644 --- a/examples/python/CuTeDSL/blackwell/mamba2_ssd/mamba2_ssd.py +++ b/examples/python/CuTeDSL/blackwell/mamba2_ssd/mamba2_ssd.py @@ -33,9 +33,6 @@ import argparse from typing import List, Type, Tuple, Optional import cuda.bindings.driver as cuda -import torch -import torch.nn.functional as F - import cutlass import cutlass.cute as cute import cutlass.cute.testing as testing @@ -43,7 +40,6 @@ import cutlass.utils as utils import cutlass.pipeline as pipeline from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait from cutlass.cute.nvgpu import cpasync, tcgen05 -import cutlass.torch as cutlass_torch import cutlass.utils.blackwell_helpers as sm100_utils from cutlass.cute.runtime import from_dlpack @@ -702,7 +698,7 @@ class SSDKernel: G = cute.size(tma_tensor_b, mode=[3]) NGROUP_RATIO = EH // G - # Make TiledMma + # Make tiledMma ( tiled_mma_intra1, tiled_mma_intra2, @@ -1670,7 +1666,10 @@ class SSDKernel: cute.copy(tiled_r2s_p, tRS_rP, tRS_sP[inter2_p_coord]) # Fence for shared memory - cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) # Async arrive INTER2_P buffer full inter2_p_pipeline.producer_commit(inter2_p_producer_state) # Advance INTER2_P producer state @@ -1700,7 +1699,10 @@ class SSDKernel: ] # Fence for shared memory - cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) # Combine B/Delta/DeltaA/last_column tScaledB = self.pre_inter_scale_bt_with_delta( @@ -1716,7 +1718,10 @@ class SSDKernel: cute.copy(tiled_r2s_b, tBrB_r2s, tBsB_r2s[inter1_b_coord]) # Fence for shared memory - cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) # Async arrive B/Delta/B_TMEM buffer empty/empty/full b_pipeline.consumer_release( @@ -1743,14 +1748,9 @@ class SSDKernel: # Combine INTER1_ACC/last_column/State exp_last_column = cute.math.exp(last_column, fastmath=True) - for reg_idx in range(0, cute.size(tTR_rP), 2): - ( - tTR_rP[reg_idx], - tTR_rP[reg_idx + 1], - ) = cute.arch.fma_packed_f32x2( - (exp_last_column, exp_last_column), - (tState[reg_idx], tState[reg_idx + 1]), - (tTR_rP[reg_idx], tTR_rP[reg_idx + 1]), + for reg_idx in cutlass.range(cute.size(tTR_rP), vectorize=True): + tTR_rP[reg_idx] = ( + exp_last_column * tState[reg_idx] + tTR_rP[reg_idx] ) # Store scaled P to tRS_rP @@ -1765,7 +1765,10 @@ class SSDKernel: cute.copy(tiled_r2s_p, tRS_rP, tRS_sP[inter2_p_coord]) # Fence for shared memory - cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) # Async arrive INTER1_ACC/INTER2_P buffer empty/full inter1_acc_pipeline.consumer_release(inter1_acc_consumer_state) @@ -1798,8 +1801,11 @@ class SSDKernel: # END of for chunk_idx in cutlass.range(C, unroll=1) # Store last INTER2_P (State) from smem to gmem - # Wait for all previous stores to smem to be done - cute.arch.fence_proxy("async.shared", space="cta") + # Wait for all previous store to smem done + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) self.pre_inter_sync_barrier.arrive_and_wait() if local_warp_idx == 0: @@ -2245,58 +2251,26 @@ class SSDKernel: cute.copy(s2r_atom_d, tRS_sD[d_coord], tRS_rD) # Combine INTRA2_ACC/INTER2_ACC/Delta/X/D - for reg_idx in range(0, cute.size(tRS_rCompute), 2): - ( - tRS_rCompute[reg_idx], - tRS_rCompute[reg_idx + 1], - ) = cute.arch.fma_packed_f32x2( - (tTR_rInter[reg_idx], tTR_rInter[reg_idx + 1]), - ( - cute.math.exp( - tTR_rDeltaA[reg_idx], fastmath=True - ), - cute.math.exp( - tTR_rDeltaA[reg_idx + 1], fastmath=True - ), - ), - (tTR_rIntra[reg_idx], tTR_rIntra[reg_idx + 1]), + for reg_idx in cutlass.range( + cute.size(tRS_rCompute), vectorize=True + ): + tRS_rCompute[reg_idx] = ( + tTR_rInter[reg_idx] + * cute.math.exp(tTR_rDeltaA[reg_idx], fastmath=True) + + tTR_rIntra[reg_idx] ) # Fuse Y += X * D if cutlass.const_expr(self.d_has_hdim): - ( - tRS_rCompute[reg_idx], - tRS_rCompute[reg_idx + 1], - ) = cute.arch.fma_packed_f32x2( - ( - tRS_rD[reg_idx].to(self.acc_dtype), - tRS_rD[reg_idx + 1].to(self.acc_dtype), - ), - ( - tSR_rX[reg_idx].to(self.acc_dtype), - tSR_rX[reg_idx + 1].to(self.acc_dtype), - ), - ( - tRS_rCompute[reg_idx], - tRS_rCompute[reg_idx + 1], - ), + tRS_rCompute[reg_idx] = ( + tRS_rD[reg_idx].to(self.acc_dtype) + * tSR_rX[reg_idx].to(self.acc_dtype) + + tRS_rCompute[reg_idx] ) elif cutlass.const_expr(self.has_d): - ( - tRS_rCompute[reg_idx], - tRS_rCompute[reg_idx + 1], - ) = cute.arch.fma_packed_f32x2( - ( - tRS_rD.to(self.acc_dtype), - tRS_rD.to(self.acc_dtype), - ), - ( - tSR_rX[reg_idx].to(self.acc_dtype), - tSR_rX[reg_idx + 1].to(self.acc_dtype), - ), - ( - tRS_rCompute[reg_idx], - tRS_rCompute[reg_idx + 1], - ), + tRS_rCompute[reg_idx] = ( + tRS_rD.to(self.acc_dtype) + * tSR_rX[reg_idx].to(self.acc_dtype) + + tRS_rCompute[reg_idx] ) tRS_rY.store(tRS_rCompute.load().to(self.io_dtype)) @@ -2309,7 +2283,10 @@ class SSDKernel: ) # Fence for R2S store - cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) # Sync before TMA store self.epilog_sync_barrier.arrive_and_wait() @@ -2426,7 +2403,6 @@ class SSDKernel: internal_stages, intra1_acc_stages, ): - SM100_TMEM_CAPACITY_COLUMNS = 512 BITS_PER_TMEM_COL = 32 # (MMA, MMA_M, MMA_N) acc_shape_intra1 = tiled_mma_intra1.partition_shape_C(tile_shape_mnk_intra1[:2]) @@ -2483,7 +2459,7 @@ class SSDKernel: num_tmem_cols_total = 1 while num_tmem_cols_total < num_tmem_cols_total_tmp: num_tmem_cols_total *= 2 - assert num_tmem_cols_total <= SM100_TMEM_CAPACITY_COLUMNS + assert num_tmem_cols_total <= cute.arch.get_max_tmem_alloc_cols("sm_100") return ( tmem_intra1_acc_offset, @@ -3036,41 +3012,26 @@ class SSDKernel: # SegSum # fadd2 + fsel + fmul2/mufu + fmul2 - for subtile_idx in cutlass.range(0, cute.size(tTR_rQ), 2, unroll_full=True): - ( - tCompute[subtile_idx], - tCompute[subtile_idx + 1], - ) = cute.arch.add_packed_f32x2( - (tCrDeltaA_Col[subtile_idx], tCrDeltaA_Col[subtile_idx + 1]), - (-tCrDeltaA_Row[subtile_idx], -tCrDeltaA_Row[subtile_idx + 1]), + for subtile_idx in cutlass.range( + cute.size(tTR_rQ), unroll_full=True, vectorize=True + ): + tCompute[subtile_idx] = tCrDeltaA_Col[subtile_idx] + ( + -tCrDeltaA_Row[subtile_idx] ) for subtile_idx in cutlass.range(cute.size(tTR_rQ), unroll_full=True): m, n = tCoord[subtile_idx] if m < n: tCompute[subtile_idx] = cutlass.Float32(-float("inf")) LOG2_E = cutlass.Float32(1.4426950408889634) - for subtile_idx in cutlass.range(0, cute.size(tTR_rQ), 2, unroll_full=True): + for subtile_idx in cutlass.range( + cute.size(tTR_rQ), unroll_full=True, vectorize=True + ): # TODO: use math.exp directly - tCompute_log2e = cute.arch.mul_packed_f32x2( - (tCompute[subtile_idx], tCompute[subtile_idx + 1]), (LOG2_E, LOG2_E) - ) - ( - tCompute[subtile_idx], - tCompute[subtile_idx + 1], - ) = cute.arch.mul_packed_f32x2( - ( - cute.math.exp2(tCompute_log2e[0], fastmath=True), - cute.math.exp2(tCompute_log2e[1], fastmath=True), - ), - (tCrDelta[subtile_idx], tCrDelta[subtile_idx + 1]), - ) - ( - tCompute[subtile_idx], - tCompute[subtile_idx + 1], - ) = cute.arch.mul_packed_f32x2( - (tCompute[subtile_idx], tCompute[subtile_idx + 1]), - (tTR_rQ[subtile_idx], tTR_rQ[subtile_idx + 1]), + tCompute_log2e = tCompute[subtile_idx] * LOG2_E + tCompute[subtile_idx] = ( + cute.math.exp2(tCompute_log2e, fastmath=True) * tCrDelta[subtile_idx] ) + tCompute[subtile_idx] = tCompute[subtile_idx] * tTR_rQ[subtile_idx] tRT_rQ.store(tCompute.load().to(self.io_dtype)) return tRT_rQ @@ -3211,6 +3172,7 @@ class SSDKernel: ) return sDeltaA + @cute.jit def pre_inter_scale_bt_with_delta( self, tBrB_s2r, tBrDelta_s2r, tBrDeltaA_s2r, last_column ): @@ -3223,22 +3185,15 @@ class SSDKernel: tBrDelta_Compute.store(tBrDelta_s2r.load().to(self.acc_dtype)) tBrDeltaA_Compute.store(tBrDeltaA_s2r.load().to(self.acc_dtype)) - for reg_idx in range(0, cute.size(tBrB_Compute), 2): - tCompute[reg_idx], tCompute[reg_idx + 1] = cute.arch.mul_packed_f32x2( - ( - cute.math.exp( - (last_column - tBrDeltaA_Compute[reg_idx]), fastmath=True - ), - cute.math.exp( - (last_column - tBrDeltaA_Compute[reg_idx + 1]), fastmath=True - ), - ), - (tBrDelta_Compute[reg_idx], tBrDelta_Compute[reg_idx + 1]), - ) - tCompute[reg_idx], tCompute[reg_idx + 1] = cute.arch.mul_packed_f32x2( - (tCompute[reg_idx], tCompute[reg_idx + 1]), - (tBrB_Compute[reg_idx], tBrB_Compute[reg_idx + 1]), + for reg_idx in cutlass.range( + cute.size(tBrB_Compute), vectorize=True, unroll_full=True + ): + tCompute[reg_idx] = ( + cute.math.exp((last_column - tBrDeltaA_Compute[reg_idx]), fastmath=True) + * tBrDelta_Compute[reg_idx] ) + + tCompute[reg_idx] = tCompute[reg_idx] * tBrB_Compute[reg_idx] return tCompute def epilog_make_delta(self, smem_cumsum_delta): @@ -3349,6 +3304,10 @@ def run( print(f"Skip reference checking: {skip_ref_check}") print(f"Use cold L2: {'True' if use_cold_l2 else 'False'}") + import torch + import torch.nn.functional as F + import cutlass.torch as cutlass_torch + # Unpack parameters G, B, E, H, C, D, L, N = gbehcdln EH = E * H diff --git a/examples/python/CuTeDSL/blackwell/mixed_input_fmha/mixed_input_fmha_decode.py b/examples/python/CuTeDSL/blackwell/mixed_input_fmha/mixed_input_fmha_decode.py index a1b14194..0fd5d365 100644 --- a/examples/python/CuTeDSL/blackwell/mixed_input_fmha/mixed_input_fmha_decode.py +++ b/examples/python/CuTeDSL/blackwell/mixed_input_fmha/mixed_input_fmha_decode.py @@ -49,53 +49,52 @@ import cutlass.torch as cutlass_torch import cutlass.utils.blackwell_helpers as sm100_utils import cutlass.cute.testing as testing from cutlass.cute.runtime import from_dlpack -from cutlass.cute.typing import Int32, Int64, Float32, Pointer, AddressSpace +from cutlass.cute.typing import * -from cutlass._mlir.dialects import nvvm, llvm +from cutlass._mlir.dialects import llvm +from cutlass.cute.arch.nvvm_wrappers import mapa # Kernel invariants mma_modes = (0, 1, 2) -mma_dice = (None, None, None) # (MMA, #MMA_M, #MMA_K) -cpy_dice = (None, *mma_dice) # (CPY, #CPY_MMA, #CPY_M, #CPY_K) +mma_dice = (None, None, None) # (MMA, #MMA_M, #MMA_K) +cpy_dice = (None,) + mma_dice # (CPY, #CPY_MMA, #CPY_M, #CPY_K) warp_threads = 32 warpgroup_warps = 4 warpgroup_threads = 128 # Math helpers -log2_e = math.log2(math.e) # change exponential base -use_tensor_ssa_math = False # experimental -fadd2 = partial(cute.arch.add_packed_f32x2, ftz=False, rnd=nvvm.FPRoundingMode.RN) -fmul2 = partial(cute.arch.mul_packed_f32x2, ftz=False, rnd=nvvm.FPRoundingMode.RN) -ffma2 = partial(cute.arch.fma_packed_f32x2, ftz=False, rnd=nvvm.FPRoundingMode.RN) +log2_e = math.log2(math.e) # change exponential base +use_tensor_ssa_math = False # experimental +fadd2 = partial(cute.arch.add_packed_f32x2, ftz=False, rnd="rn") +fmul2 = partial(cute.arch.mul_packed_f32x2, ftz=False, rnd="rn") +ffma2 = partial(cute.arch.fma_packed_f32x2, ftz=False, rnd="rn") exp2 = partial(cute.math.exp2, fastmath=True) + class MixedInputFusedMultiHeadAttentionDecode: def __init__( self, headdim, block_scaledim, # headdim per scale factor; scale factor shape is (batches, heads_k, seqlen, headdim / block_scaledim) grouped_head_tile, # GQA packing tile size, can be less than group size - dual_convert = False, # Dual warpgroups pingponging on convert stages - deterministic = False, # If True, cluster reduction is disabled + convert_warpgroups = 1, # Multiple warpgroups striding on convert stages ): self.headdim = headdim self.grouped_head_tile = grouped_head_tile self.block_scaledim = block_scaledim self.scaledim = headdim // block_scaledim - self.dual_convert = dual_convert - self.deterministic = deterministic + self.convert_warpgroups = convert_warpgroups assert headdim % block_scaledim == 0 assert grouped_head_tile % 8 == 0 and 0 < grouped_head_tile <= 32 - + warpgroup_id = 0 self.softmax_warpgroup_id = warpgroup_id warpgroup_id += 1 - self.cvt_warpgroup_ids = (warpgroup_id, (warpgroup_id+1) if dual_convert else None) - cvt_warpgroups = 2 if dual_convert else 1 - warpgroup_id += cvt_warpgroups + self.cvt_warpgroup_ids = tuple(range(warpgroup_id, warpgroup_id+convert_warpgroups)) + warpgroup_id += convert_warpgroups # Why 2 MMA+TMA warps when not MMA bound? # Less register pressure per warp promotes concise SASS @@ -112,53 +111,65 @@ class MixedInputFusedMultiHeadAttentionDecode: self.threads_per_cta = warpgroup_id * warpgroup_threads self.use_reg_reconfig = grouped_head_tile > 16 - max_regs_per_wg_thread = (64*1024 // warpgroup_threads) # 64K regs per SM + max_regs_per_wg_thread = 64 * 1024 // warpgroup_threads # 64K regs per SM self.mma_tma_regs = 72 self.cvt_regs = 112 - self.softmax_regs = min(256, max_regs_per_wg_thread - self.mma_tma_regs - - self.cvt_regs * cvt_warpgroups) + self.softmax_regs = (max_regs_per_wg_thread + - self.mma_tma_regs + - self.cvt_regs * convert_warpgroups) + self.softmax_regs = max(128, min(256, self.softmax_regs)) assert (self.mma_tma_regs + self.softmax_regs + - self.cvt_regs * cvt_warpgroups) <= max_regs_per_wg_thread + self.cvt_regs * convert_warpgroups) <= max_regs_per_wg_thread or not self.use_reg_reconfig self.bs_stages = 2 self.sp_stages = 2 self.o_stages = 1 - def can_implement(self, problem_shape, kv_splits, kv_cluster_dim, q_dtype, kv_dtype, o_dtype, acc_dtype): + def can_implement( + self, + problem_shape, + kv_splits, + q_dtype, + kv_dtype, + o_dtype, + acc_dtype, + ): b, h_q, h_k, s_k, d = problem_shape + if kv_dtype is cutlass.Float8E4M3: + raise ValueError("use Float8E4M3FN instead of Float8E4M3") + + if d % 64 != 0: + raise ValueError(f"headdim({d}) must be multiple of 64") + if h_q % h_k != 0: raise ValueError(f"heads_q({h_q}) must be a multiple of heads_k({h_k})") - - if kv_splits % kv_cluster_dim != 0: - raise ValueError(f"kv_splits({kv_splits}) must be a multiple of kv_cluster_dim({kv_cluster_dim})") - if self.deterministic and kv_cluster_dim != 1: - raise ValueError(f"kv_cluster_dim({kv_cluster_dim}) must be 1 for determinism") - - align_scale_bits = 128 # TMA requirement + align_scale_bits = 128 # TMA requirement if self.scaledim * q_dtype.width < align_scale_bits: align_seq = align_scale_bits // (self.scaledim * q_dtype.width) if s_k % align_seq != 0: raise ValueError(f"seqlen({s_k}) must be a multiple of {align_seq}") + if kv_dtype.width < 8 and d % 128 != 0: # TMA requirement + raise ValueError(f"headdim({d}) must be multiple of 128 for {kv_dtype} KV") + @cute.jit def __call__( self, - problem_shape: Tuple[Int32, Int32, Int32, Int32, Int32], # b, h_q, h_k, s_k, d - kv_splits: Int32, # threadblocks per sequence - kv_cluster_dim: Int32, # threadblocks per partial buffer (atomic reduction) + problem_shape: Tuple[Int32, Int32, Int32, Int32, Int32], # b, h_q, h_k, s_k, d + kv_splits: Int32, # threadblocks per sequence q_iter: cute.Pointer, k_iter: cute.Pointer, v_iter: cute.Pointer, k_scale_iter: cute.Pointer, v_scale_iter: cute.Pointer, o_iter: cute.Pointer, - m_iter: cute.Pointer, # colmax_s, must be -inf initialized - l_iter: cute.Pointer, # logsumexp - o_partial_iter: cute.Pointer, # partial O per kv cluster, must be zero initialized if nondeterminism is enabled - m_partial_iter: cute.Pointer, # partial colmax_s per kv cluster - l_partial_iter: cute.Pointer, # partial colsum_p per kv cluster, must be zero initialized if nondeterminism is enabled + m_iter: cute.Pointer, # colmax_s, must be -inf initialized + l_iter: cute.Pointer, # logsumexp + o_partial_iter: cute.Pointer, # partial O per kv split + m_partial_iter: cute.Pointer, # partial colmax_s per kv split + l_partial_iter: cute.Pointer, # partial colsum_p per kv split scale_qs: Float32, scale_o: Float32, stream: cuda.CUstream, @@ -168,7 +179,7 @@ class MixedInputFusedMultiHeadAttentionDecode: ############################## mma_dtype = q_iter.dtype acc_dtype = o_partial_iter.dtype - assert acc_dtype is Float32 # don't support other acc types for now + assert acc_dtype is Float32 # don't support other acc types for now # Block tile sets the granularity at which threadblocks consume work blk_tile_s = 128 @@ -186,29 +197,30 @@ class MixedInputFusedMultiHeadAttentionDecode: # GEMM1: (S_K, H_R, D, (H_K, B)) tiled_mma_kq = sm100_utils.make_trivial_tiled_mma( mma_dtype, - tcgen05.OperandMajorMode.K, # K - tcgen05.OperandMajorMode.K, # Q + tcgen05.OperandMajorMode.K, # K + tcgen05.OperandMajorMode.K, # Q acc_dtype, tcgen05.CtaGroup.ONE, mma_tile_mnk[:2], - tcgen05.OperandSource.TMEM, # converted K in tmem + tcgen05.OperandSource.TMEM, # converted K in tmem ) # GEMM2: (D, H_R, S_K, (H_K, B)) - tiled_mma_vp = sm100_utils.make_trivial_tiled_mma(# + tiled_mma_vp = sm100_utils.make_trivial_tiled_mma( # mma_dtype, - tcgen05.OperandMajorMode.K, # V - tcgen05.OperandMajorMode.MN, # P + tcgen05.OperandMajorMode.K, # V + tcgen05.OperandMajorMode.MN, # P acc_dtype, tcgen05.CtaGroup.ONE, mma_tile_mnk[:2], - tcgen05.OperandSource.TMEM, # converted V in tmem + tcgen05.OperandSource.TMEM, # converted V in tmem ) # Calculate Q stages self.q_stages = blk_tile_d // mma_tile_k - # Heuristics to avoid power throttling + # Perf heuristics + cap_kv_stages = k_iter.dtype.width >= 8 max_cvt_stages = 4 if self.grouped_head_tile == 32 and mma_tile_k == 128 else 8 max_kv_stages = 8 if mma_tile_k == 128 else 14 @@ -218,25 +230,40 @@ class MixedInputFusedMultiHeadAttentionDecode: tmem_capacity = 512 cvt_stage_cols = mma_tile_k * mma_dtype.width // 32 self.cvt_stages = (tmem_capacity - tmem_alloc_cols) // cvt_stage_cols - self.cvt_stages = min(self.cvt_stages, max_cvt_stages) + self.cvt_stages = ( + min(self.cvt_stages, max_cvt_stages) if cap_kv_stages else self.cvt_stages + ) tmem_alloc_cols += self.cvt_stages * cvt_stage_cols - self.tmem_alloc_cols = 2 ** math.ceil(math.log2(tmem_alloc_cols)) # Tmem alloc must be PO2 + self.tmem_alloc_cols = 2 ** math.ceil( + math.log2(tmem_alloc_cols) + ) # Tmem alloc must be PO2 print(f"\tcvt stages: {self.cvt_stages}") # Calculate KV smem stages self.mbarrier_reserved_bytes = 768 - smem_alloc_bits = self.mbarrier_reserved_bytes*8 - smem_alloc_bits += mma_tile_n * 2 * acc_dtype.width # colmax + cluster colmax - smem_alloc_bits += self.scaledim * blk_tile_s * self.bs_stages * mma_dtype.width # block scale - smem_alloc_bits += mma_tile_n * warpgroup_warps * acc_dtype.width # colsum - smem_alloc_bits += mma_tile_n * mma_tile_k * self.q_stages * mma_dtype.width # Q - smem_alloc_bits += mma_tile_m * mma_tile_n * self.sp_stages * mma_dtype.width # P + smem_alloc_bits = self.mbarrier_reserved_bytes * 8 + smem_alloc_bits += mma_tile_n * acc_dtype.width # colmax + smem_alloc_bits += ( + self.scaledim * blk_tile_s * self.bs_stages * mma_dtype.width + ) # block scale + smem_alloc_bits += mma_tile_n * warpgroup_warps * acc_dtype.width # colsum + smem_alloc_bits += ( + mma_tile_n * mma_tile_k * self.q_stages * mma_dtype.width + ) # Q + smem_alloc_bits += ( + mma_tile_m * mma_tile_n * self.sp_stages * mma_dtype.width + ) # P smem_capacity = utils.get_smem_capacity_in_bytes("sm_100") - self.kv_stages = (smem_capacity*8 - smem_alloc_bits) // (mma_tile_m * mma_tile_k * k_iter.dtype.width) - self.kv_stages = min(self.kv_stages, max_kv_stages) + kv_smem_dtype = cutlass.Int8 if k_iter.dtype.width < 8 else k_iter.dtype + self.kv_stages = (smem_capacity * 8 - smem_alloc_bits) // ( + mma_tile_m * mma_tile_k * kv_smem_dtype.width + ) + self.kv_stages = ( + min(self.kv_stages, max_kv_stages) if cap_kv_stages else self.kv_stages + ) print(f"\tkv stages: {self.kv_stages}") @@ -245,74 +272,68 @@ class MixedInputFusedMultiHeadAttentionDecode: ############################## b, h_q, h_k, s_k, d = problem_shape h_r = h_q // h_k - kv_clusters = kv_splits // kv_cluster_dim - q = cute.make_tensor(q_iter, - cute.make_ordered_layout( - shape=(h_r, d, (h_k, b)), - order=( 1, 0, ( 2, 3)) - ) + q = cute.make_tensor( + q_iter, + cute.make_ordered_layout(shape=(h_r, d, (h_k, b)), order=(1, 0, (2, 3))), ) - k = cute.make_tensor(k_iter, - cute.make_ordered_layout( - shape=(s_k, d, (h_k, b)), - order=( 1, 0, ( 2, 3)) - ) + k = cute.make_tensor( + k_iter, + cute.make_ordered_layout(shape=(s_k, d, (h_k, b)), order=(1, 0, (2, 3))), ) assert k_iter.dtype is not q_iter.dtype - v = cute.make_tensor(v_iter, - cute.make_ordered_layout( - shape=(d, s_k, (h_k, b)), - order=(0, 1, ( 2, 3)) - ) + v = cute.make_tensor( + v_iter, + cute.make_ordered_layout(shape=(d, s_k, (h_k, b)), order=(0, 1, (2, 3))), ) assert v_iter.dtype is k_iter.dtype - o_partial = cute.make_tensor(o_partial_iter, + o_partial = cute.make_tensor( + o_partial_iter, cute.make_ordered_layout( - shape=(d, h_r, (h_k, b), kv_clusters), - order=(0, 1, ( 2, 3), 4) - ) + shape=(d, h_r, (h_k, b), kv_splits), order=(0, 1, (2, 3), 4) + ), ) - m = cute.make_tensor(m_iter, + m = cute.make_tensor( + m_iter, cute.make_ordered_layout( shape=(h_r, (h_k, b)), - order=( 0, ( 1, 2)), - ) + order=(0, (1, 2)), + ), ) assert m_iter.dtype is acc_dtype - m_partial = cute.make_tensor(m_partial_iter, + m_partial = cute.make_tensor( + m_partial_iter, cute.make_ordered_layout( - shape=(h_r, (h_k, b), kv_clusters), - order=( 0, ( 1, 2), 3), - ) + shape=(h_r, (h_k, b), kv_splits), + order=(0, (1, 2), 3), + ), ) assert m_partial_iter.dtype is acc_dtype - l_partial = cute.make_tensor(l_partial_iter, + l_partial = cute.make_tensor( + l_partial_iter, cute.make_ordered_layout( - shape=(h_r, (h_k, b), kv_clusters), - order=( 0, ( 1, 2), 3), - ) + shape=(h_r, (h_k, b), kv_splits), + order=(0, (1, 2), 3), + ), ) assert l_partial_iter.dtype is acc_dtype - align_scale_bits = 128 # TMA requirement + align_scale_bits = 128 # TMA requirement if cutlass.const_expr(self.scaledim * mma_dtype.width >= align_scale_bits): scale_layout = cute.make_ordered_layout( - shape=(self.scaledim, s_k, (h_k, b)), - order=(0, 1, (2, 3)) + shape=(self.scaledim, s_k, (h_k, b)), order=(0, 1, (2, 3)) ) else: align_seq = align_scale_bits // (self.scaledim * mma_dtype.width) s_ks = (align_seq, s_k // align_seq) scale_layout = cute.make_ordered_layout( - shape=(self.scaledim, s_ks, (h_k, b)), - order=(0, (1, 2), (3, 4)) + shape=(self.scaledim, s_ks, (h_k, b)), order=(0, (1, 2), (3, 4)) ) k_scale = cute.make_tensor(k_scale_iter, scale_layout) @@ -321,40 +342,70 @@ class MixedInputFusedMultiHeadAttentionDecode: v_scale = cute.make_tensor(v_scale_iter, scale_layout) assert v_scale_iter.dtype is mma_dtype - # (MMA, MMA_M/N, MMA_K, Stages) - smem_layout_q = sm100_utils.make_smem_layout_b(tiled_mma_kq, mma_tile_mnk, q_iter.dtype, self.q_stages) - smem_layout_k = sm100_utils.make_smem_layout_a(tiled_mma_kq, mma_tile_mnk, k_iter.dtype, self.kv_stages) - smem_layout_v = sm100_utils.make_smem_layout_a(tiled_mma_vp, mma_tile_mnk, v_iter.dtype, self.kv_stages, - is_k_major=False) # V is always headdim-major (GEMM2 M-major) in gmem+smem + smem_layout_q = sm100_utils.make_smem_layout_b( + tiled_mma_kq, mma_tile_mnk, q_iter.dtype, self.q_stages + ) + smem_layout_k = sm100_utils.make_smem_layout_a( + tiled_mma_kq, mma_tile_mnk, kv_smem_dtype, self.kv_stages + ) + smem_layout_v = sm100_utils.make_smem_layout_a( + tiled_mma_vp, mma_tile_mnk, kv_smem_dtype, self.kv_stages, is_k_major=False + ) # V is always headdim-major (GEMM2 M-major) in gmem+smem smem_layout_bs = cute.make_layout((self.scaledim, blk_tile_s, self.bs_stages)) - smem_layout_atom_o = tcgen05.make_smem_layout_atom(tcgen05.mma.SmemLayoutAtomKind.MN_SW128, o_partial_iter.dtype) - smem_layout_o = cute.tile_to_shape(smem_layout_atom_o, (blk_tile_d, blk_tile_h), order=(1,0)) + smem_layout_atom_o = tcgen05.make_smem_layout_atom( + tcgen05.mma.SmemLayoutAtomKind.MN_SW128, o_partial_iter.dtype + ) + smem_layout_o = cute.tile_to_shape( + smem_layout_atom_o, (blk_tile_d, blk_tile_h), order=(1, 0) + ) smem_layout_o = cute.flat_divide(smem_layout_o, (mma_tile_m, mma_tile_n)) tma_load_op = cute.nvgpu.cpasync.CopyBulkTensorTileG2SOp() - tma_store_op = (cute.nvgpu.cpasync.CopyBulkTensorTileS2GOp() if self.deterministic - else cute.nvgpu.cpasync.CopyReduceBulkTensorTileS2GOp()) + tma_store_op = cute.nvgpu.cpasync.CopyBulkTensorTileS2GOp() tma_atom_q, tma_tensor_q = cute.nvgpu.make_tiled_tma_atom_B( - tma_load_op, q, cute.select(smem_layout_q, mma_modes), mma_tile_mnk, tiled_mma_kq + tma_load_op, + q, + cute.select(smem_layout_q, mma_modes), + mma_tile_mnk, + tiled_mma_kq, ) tma_atom_k, tma_tensor_k = cute.nvgpu.make_tiled_tma_atom_A( - tma_load_op, k, cute.select(smem_layout_k, mma_modes), mma_tile_mnk, tiled_mma_kq + tma_load_op, + k, + cute.select(smem_layout_k, mma_modes), + mma_tile_mnk, + tiled_mma_kq, + internal_type=kv_smem_dtype, ) tma_atom_v, tma_tensor_v = cute.nvgpu.make_tiled_tma_atom_A( - tma_load_op, v, cute.select(smem_layout_v, mma_modes), mma_tile_mnk, tiled_mma_vp + tma_load_op, + v, + cute.select(smem_layout_v, mma_modes), + mma_tile_mnk, + tiled_mma_vp, + internal_type=kv_smem_dtype, ) tma_atom_ks, tma_tensor_ks = cute.nvgpu.cpasync.make_tiled_tma_atom( - tma_load_op, k_scale, cute.select(smem_layout_bs, mode=[0,1]), smem_layout_bs.shape[:2] + tma_load_op, + k_scale, + cute.select(smem_layout_bs, mode=[0, 1]), + smem_layout_bs.shape[:2], ) tma_atom_vs, tma_tensor_vs = cute.nvgpu.cpasync.make_tiled_tma_atom( - tma_load_op, v_scale, cute.select(smem_layout_bs, mode=[0,1]), smem_layout_bs.shape[:2] + tma_load_op, + v_scale, + cute.select(smem_layout_bs, mode=[0, 1]), + smem_layout_bs.shape[:2], ) tma_atom_o, tma_tensor_o = cute.nvgpu.cpasync.make_tiled_tma_atom( - tma_store_op, o_partial, cute.select(smem_layout_o, mode=[0,1]), mma_tile_mnk[:2] + tma_store_op, + o_partial, + cute.select(smem_layout_o, mode=[0, 1]), + mma_tile_mnk[:2], ) # K scale and V scale will have the same TMA tensor (coord tensor) @@ -371,18 +422,39 @@ class MixedInputFusedMultiHeadAttentionDecode: grid = (kv_splits, n_tiles, l_tiles) self.decode( - blk_tile_shd, mma_tile_mnk, tiled_mma_kq, tiled_mma_vp, - q_iter.dtype, smem_layout_q, tma_atom_q, tma_tensor_q, - k_iter.dtype, smem_layout_k, tma_atom_k, tma_tensor_k, - v_iter.dtype, smem_layout_v, tma_atom_v, tma_tensor_v, - smem_layout_bs, tma_atom_ks, tma_atom_vs, tma_tensor_bs, - o_partial_iter.dtype, smem_layout_o, tma_atom_o, tma_tensor_o, - m, m_partial, l_partial, - scale_qs, scale_qs_log2_e, + blk_tile_shd, + mma_tile_mnk, + tiled_mma_kq, + tiled_mma_vp, + q_iter.dtype, + smem_layout_q, + tma_atom_q, + tma_tensor_q, + k_iter.dtype, + smem_layout_k, + tma_atom_k, + tma_tensor_k, + v_iter.dtype, + smem_layout_v, + tma_atom_v, + tma_tensor_v, + smem_layout_bs, + tma_atom_ks, + tma_atom_vs, + tma_tensor_bs, + o_partial_iter.dtype, + smem_layout_o, + tma_atom_o, + tma_tensor_o, + m, + m_partial, + l_partial, + scale_qs, + scale_qs_log2_e, ).launch( grid=grid, block=[self.threads_per_cta, 1, 1], - cluster=[kv_cluster_dim, 1, 1], + cluster=[1, 1, 1], stream=stream, min_blocks_per_mp=1, ) @@ -394,9 +466,15 @@ class MixedInputFusedMultiHeadAttentionDecode: m = cute.make_tensor(m_iter, cute.make_layout((h_q, b))) l = cute.make_tensor(l_iter, cute.make_layout((h_q, b))) - o_partial = cute.make_tensor(o_partial_iter, cute.make_layout((d, h_q, b, kv_clusters))) - m_partial = cute.make_tensor(m_partial_iter, cute.make_layout((h_q, b, kv_clusters))) - l_partial = cute.make_tensor(l_partial_iter, cute.make_layout((h_q, b, kv_clusters))) + o_partial = cute.make_tensor( + o_partial_iter, cute.make_layout((d, h_q, b, kv_splits)) + ) + m_partial = cute.make_tensor( + m_partial_iter, cute.make_layout((h_q, b, kv_splits)) + ) + l_partial = cute.make_tensor( + l_partial_iter, cute.make_layout((h_q, b, kv_splits)) + ) d_per_blk = 128 d_blks = cute.ceil_div(d, d_per_blk) @@ -451,9 +529,6 @@ class MixedInputFusedMultiHeadAttentionDecode: # Read special registers kv_splits, tiles_hr, tiles_hb = cute.arch.grid_dim() kv_split_idx, coord_hr, coord_hb = cute.arch.block_idx() - kv_split_in_cluster, _, _ = cute.arch.block_in_cluster_idx() - kv_cluster_dim, _, _ = cute.arch.block_in_cluster_dim() - kv_cluster_idx, _, _ = cute.arch.cluster_idx() tidx, _, _ = cute.arch.thread_idx() lane_idx = cute.arch.lane_idx() warp_idx = cute.arch.make_warp_uniform(tidx // warp_threads) @@ -464,17 +539,22 @@ class MixedInputFusedMultiHeadAttentionDecode: # No multicast mcast_coord = 0 - mcast_layout = cute.make_layout((1,1,1,1)) # vmnk + mcast_layout = cute.make_layout((1, 1, 1, 1)) # vmnk # Alias types mma_dtype = q_dtype acc_dtype = o_dtype + kv_smem_dtype = cutlass.Int8 if k_dtype.width < 8 else k_dtype # Shapes for MMA tile indexing (Read TMA partition for example) blk_tile_s, blk_tile_h, blk_tile_d = blk_tile_shd mma_tile_m, mma_tile_n, mma_tile_k = mma_tile_mnk - tiles_dm, tiles_sk = cute.ceil_div((blk_tile_d, blk_tile_s), (mma_tile_m, mma_tile_k)) - tiles_dk, tiles_sm = cute.ceil_div((blk_tile_d, blk_tile_s), (mma_tile_k, mma_tile_m)) + tiles_dm, tiles_sk = cute.ceil_div( + (blk_tile_d, blk_tile_s), (mma_tile_m, mma_tile_k) + ) + tiles_dk, tiles_sm = cute.ceil_div( + (blk_tile_d, blk_tile_s), (mma_tile_k, mma_tile_m) + ) tiles_s = cute.ceil_div(mK.shape[0], blk_tile_s) iters_s = cute.ceil_div(tiles_s - kv_split_idx, kv_splits) prefetch_iters = self.sp_stages - 1 @@ -484,7 +564,6 @@ class MixedInputFusedMultiHeadAttentionDecode: # Runtime checks exit_early = kv_split_idx >= tiles_s - do_cluster_reduction = not self.deterministic and kv_cluster_dim > 1 lane_store_max = mma_tile_n == warp_threads or lane_idx < mma_tile_n # Smem alloc helper @@ -523,7 +602,6 @@ class MixedInputFusedMultiHeadAttentionDecode: s_pipeline_ptr = smem.allocate_array(Int64, self.sp_stages * 2) p_pipeline_ptr = smem.allocate_array(Int64, self.sp_stages * 2) o_pipeline_ptr = smem.allocate_array(Int64, self.o_stages * 2) - m_cluster_full_ptr = smem.allocate_array(Int64) # signal cluster colmax is in split 0 smem assert smem._allocated_bytes <= self.mbarrier_reserved_bytes @@ -532,92 +610,75 @@ class MixedInputFusedMultiHeadAttentionDecode: mma_kq_nbar_id = 2 mma_vp_nbar_id = 3 - # Initialize cluster colmax + mbar (even if this split exits early) - sM_layout = cute.make_layout(shape=(mma_tile_m, mma_tile_n), stride=(0, 1)) - sM_cluster = smem.allocate_tensor(acc_dtype, sM_layout, svector_align) - if do_cluster_reduction: - if warp_idx == init_warp and kv_split_in_cluster == 0: - if lane_store_max: - sM_cluster[(0,lane_idx)] = -Float32.inf - cute.arch.fence_acq_rel_cluster() - init_warp += 1 - - if warp_idx == init_warp: - # split 0 waits for one arrive per colmax elt per split in cluster - # other splits in cluster wait for one arrive from split 0 - arrive_count = (kv_cluster_dim * mma_tile_n) if kv_split_in_cluster == 0 else 1 - cute.arch.mbarrier_init(m_cluster_full_ptr, arrive_count) - cute.arch.mbarrier_init_fence() - init_warp += 1 - - cute.arch.cluster_arrive_relaxed() - - # Setup up thread cooperatives + # Alias thread cooperatives elect_one_cooperative = pipeline.CooperativeGroup(pipeline.Agent.Thread) warpgroup_cooperative = pipeline.CooperativeGroup(pipeline.Agent.Thread, warpgroup_threads) - dual_warpgroup_cooperative = pipeline.CooperativeGroup(pipeline.Agent.Thread, warpgroup_threads * 2) mma_group = elect_one_cooperative tma_group = elect_one_cooperative - tma_qo_group = elect_one_cooperative cvt_group = warpgroup_cooperative + cvt_groups = pipeline.CooperativeGroup(pipeline.Agent.Thread, warpgroup_threads * self.convert_warpgroups) softmax_group = warpgroup_cooperative # Initialize pipelines q_producer, q_consumer = pipeline.PipelineTmaAsync.create( num_stages=self.q_stages, producer_group=tma_group, - consumer_group=softmax_group, # Reuse Q consumer mbarriers to sync O store + consumer_group=softmax_group, # Reuse Q consumer mbarriers to sync O store tx_count=cute.size_in_bytes(q_dtype, cute.select(smem_layout_q, mma_modes)), barrier_storage=q_pipeline_ptr, - tidx=mcast_coord, cta_layout_vmnk=mcast_layout, + tidx=mcast_coord, + cta_layout_vmnk=mcast_layout, defer_sync=True, ).make_participants() kv_producer, kv_consumer = pipeline.PipelineTmaAsync.create( num_stages=self.kv_stages, - producer_group=tma_group, consumer_group=cvt_group, + producer_group=tma_group, + consumer_group=cvt_group, tx_count=cute.size_in_bytes(k_dtype, cute.select(smem_layout_k, mma_modes)), barrier_storage=kv_pipeline_ptr, - tidx=mcast_coord, cta_layout_vmnk=mcast_layout, + tidx=mcast_coord, + cta_layout_vmnk=mcast_layout, defer_sync=True, ).make_participants() bs_producer, bs_consumer = pipeline.PipelineTmaAsync.create( num_stages=self.bs_stages, - producer_group=tma_qo_group, - consumer_group=(dual_warpgroup_cooperative if self.dual_convert else cvt_group), + producer_group=tma_group, + consumer_group=cvt_groups, tx_count=cute.size_in_bytes(mma_dtype, cute.select(smem_layout_bs, mode=[0,1])), barrier_storage=bs_pipeline_ptr, - tidx=mcast_coord, cta_layout_vmnk=mcast_layout, + tidx=mcast_coord, + cta_layout_vmnk=mcast_layout, defer_sync=True, ).make_participants() cvt_producer, cvt_consumer = pipeline.PipelineAsyncUmma.create( num_stages=self.cvt_stages, - producer_group=cvt_group, consumer_group=mma_group, + producer_group=cvt_group, + consumer_group=mma_group, barrier_storage=cvt_pipeline_ptr, defer_sync=True, ).make_participants() s_producer, s_consumer = pipeline.PipelineUmmaAsync.create( num_stages=self.sp_stages, - producer_group=mma_group, consumer_group=softmax_group, + producer_group=mma_group, + consumer_group=softmax_group, barrier_storage=s_pipeline_ptr, defer_sync=True, ).make_participants() p_producer, p_consumer = pipeline.PipelineAsyncUmma.create( num_stages=self.sp_stages, - producer_group=softmax_group, consumer_group=mma_group, + producer_group=softmax_group, + consumer_group=mma_group, barrier_storage=p_pipeline_ptr, defer_sync=True, ).make_participants() o_producer, o_consumer = pipeline.PipelineUmmaAsync.create( num_stages=self.o_stages, - producer_group=mma_group, consumer_group=softmax_group, + producer_group=mma_group, + consumer_group=softmax_group, barrier_storage=o_pipeline_ptr, defer_sync=True, ).make_participants() - # Ensure visibility of cluster mbarrier + colmax init - if do_cluster_reduction: - cute.arch.cluster_wait() - # Ensure visibility of local mbarrier inits and tmem alloc cute.arch.sync_threads() @@ -629,41 +690,65 @@ class MixedInputFusedMultiHeadAttentionDecode: thrblk_mma_vp = tiled_mma_vp.get_slice(0) # M - colmax + sM_layout = cute.make_layout(shape=(mma_tile_m, mma_tile_n), stride=(0, 1)) sM = smem.allocate_tensor(acc_dtype, sM_layout, svector_align) tCsM = thrblk_mma_kq.partition_C(sM) - tCsM_cluster = thrblk_mma_kq.partition_C(sM_cluster) # L - colsum - sL_layout = cute.make_layout(shape=(mma_tile_m, mma_tile_n, warpgroup_warps), stride=(0, 1, mma_tile_n)) + sL_layout = cute.make_layout( + shape=(mma_tile_m, mma_tile_n, warpgroup_warps), stride=(0, 1, mma_tile_n) + ) sL = smem.allocate_tensor(acc_dtype, sL_layout, svector_align) tCsL = thrblk_mma_kq.partition_C(sL) # BS - block scale - sBS = smem.allocate_tensor(mma_dtype, smem_layout_bs, stensor_align) # (SCALE, TILE_S, bs_stages) + sBS = smem.allocate_tensor( + mma_dtype, smem_layout_bs, stensor_align + ) # (SCALE, TILE_S, bs_stages) # Q - tBsQ = smem.allocate_tensor(q_dtype, smem_layout_q.outer, stensor_align, smem_layout_q.inner) # (MMA, #MMA_N, #MMA_K, q_stages) + tBsQ = smem.allocate_tensor( + q_dtype, smem_layout_q.outer, stensor_align, smem_layout_q.inner + ) # (MMA, #MMA_N, #MMA_K, q_stages) # K - tAsK = smem.allocate_tensor(k_dtype, smem_layout_k.outer, stensor_align, smem_layout_k.inner) # (MMA, #MMA_M, #MMA_K, kv_stages) - tAtK_cvt_shape = tiled_mma_kq.partition_shape_A((mma_tile_m, mma_tile_k, self.cvt_stages)) # (MMA, #MMA_M, #MMA_K, cvt_stages) + tAsK = smem.allocate_tensor( + kv_smem_dtype, smem_layout_k.outer, stensor_align, smem_layout_k.inner + ) # (MMA, #MMA_M, #MMA_K, kv_stages) + tAtK_cvt_shape = tiled_mma_kq.partition_shape_A( + (mma_tile_m, mma_tile_k, self.cvt_stages) + ) # (MMA, #MMA_M, #MMA_K, cvt_stages) tAtK_cvt = thrblk_mma_kq.make_fragment_A(tAtK_cvt_shape) # V - tAsV_iterator = cute.recast_ptr(tAsK.iterator, smem_layout_v.inner, dtype=v_dtype) # KV share input buffers - tAsV = cute.make_tensor(tAsV_iterator, smem_layout_v.outer) # (MMA, #MMA_M, #MMA_K, kv_stages) - tAtV_cvt_shape = tiled_mma_vp.partition_shape_A((mma_tile_m, mma_tile_k, self.cvt_stages)) # (MMA, #MMA_M, #MMA_K, cvt_stages) + tAsV_iterator = cute.recast_ptr( + tAsK.iterator, smem_layout_v.inner, dtype=kv_smem_dtype + ) # KV share input buffers + tAsV = cute.make_tensor( + tAsV_iterator, smem_layout_v.outer + ) # (MMA, #MMA_M, #MMA_K, kv_stages) + tAtV_cvt_shape = tiled_mma_vp.partition_shape_A( + (mma_tile_m, mma_tile_k, self.cvt_stages) + ) # (MMA, #MMA_M, #MMA_K, cvt_stages) tAtV_cvt = thrblk_mma_vp.make_fragment_A(tAtV_cvt_shape) # S - tCtS_shape = tiled_mma_kq.partition_shape_C((mma_tile_m, mma_tile_n, self.sp_stages)) - tCtS = thrblk_mma_kq.make_fragment_C(tCtS_shape) # (MMA_MN, #MMA_M=1, #MMA_N=1, sp_stages) + tCtS_shape = tiled_mma_kq.partition_shape_C( + (mma_tile_m, mma_tile_n, self.sp_stages) + ) + tCtS = thrblk_mma_kq.make_fragment_C( + tCtS_shape + ) # (MMA_MN, #MMA_M=1, #MMA_N=1, sp_stages) # P - Treat MN C tile of BMM0 as NM B tile of BMM1 # (MMA_NK, #MMA_N, #MMA_K=MMA_TILE_M/MMA_K, sp_stages) mma_tile_nm = (None, mma_tile_n, mma_tile_m) - tBsP_nm_layout = sm100_utils.make_smem_layout_b(tiled_mma_vp, mma_tile_nm, mma_dtype, self.sp_stages) - tBsP_nm = smem.allocate_tensor(mma_dtype, tBsP_nm_layout.outer, stensor_align, tBsP_nm_layout.inner) + tBsP_nm_layout = sm100_utils.make_smem_layout_b( + tiled_mma_vp, mma_tile_nm, mma_dtype, self.sp_stages + ) + tBsP_nm = smem.allocate_tensor( + mma_dtype, tBsP_nm_layout.outer, stensor_align, tBsP_nm_layout.inner + ) # Tile for NK B tile iteration # (MMA_NK, #MMA_N, #MMA_K=MMA_TILE_K/MMA_K, #TILES_SK=MMA_TILE_M/MMA_TILE_K, sp_stages) @@ -677,39 +762,49 @@ class MixedInputFusedMultiHeadAttentionDecode: tCsP = cute.composition(tBsP_nm, tCsP_tile) # O - sO_iterator = cute.recast_ptr(tBsQ.iterator, smem_layout_o.inner, dtype=o_dtype) # Reuse QKV smem for O TMA store - sO_mma = cute.make_tensor(sO_iterator, smem_layout_o.outer) # (MMA_TILE_M, MMA_TILE_N, #TILE_DM, #TILE_HN) - tCsO = thrblk_mma_vp.partition_C(sO_mma) # (MMA, #MMA_M, #MMA_N, #TILE_DM, #TILE_HN) + sO_iterator = cute.recast_ptr( + tBsQ.iterator, smem_layout_o.inner, dtype=o_dtype + ) # Reuse QKV smem for O TMA store + sO_mma = cute.make_tensor( + sO_iterator, smem_layout_o.outer + ) # (MMA_TILE_M, MMA_TILE_N, #TILE_DM, #TILE_HN) + tCsO = thrblk_mma_vp.partition_C( + sO_mma + ) # (MMA, #MMA_M, #MMA_N, #TILE_DM, #TILE_HN) tCtO = thrblk_mma_vp.make_fragment_C(tCsO.shape) # Tmem tensor allocation tmem_ptr = cute.arch.retrieve_tmem_ptr(Int32, 16, tmem_ptr_smem_ptr) tmem_offset = 0 - tAtK_cvt = cute.make_tensor(cute.recast_ptr(tmem_ptr + tmem_offset, dtype=mma_dtype), tAtK_cvt.layout) - tAtV_cvt = cute.make_tensor(cute.recast_ptr(tmem_ptr + tmem_offset, dtype=mma_dtype), tAtV_cvt.layout) + tAtK_cvt = cute.make_tensor( + cute.recast_ptr(tmem_ptr + tmem_offset, dtype=mma_dtype), tAtK_cvt.layout + ) + tAtV_cvt = cute.make_tensor( + cute.recast_ptr(tmem_ptr + tmem_offset, dtype=mma_dtype), tAtV_cvt.layout + ) tmem_offset += tcgen05.find_tmem_tensor_col_offset(tAtK_cvt) - tCtS = cute.make_tensor(cute.recast_ptr(tmem_ptr + tmem_offset, dtype=acc_dtype), tCtS.layout) + tCtS = cute.make_tensor( + cute.recast_ptr(tmem_ptr + tmem_offset, dtype=acc_dtype), tCtS.layout + ) tmem_offset += tcgen05.find_tmem_tensor_col_offset(tCtS) - tCtO = cute.make_tensor(cute.recast_ptr(tmem_ptr + tmem_offset, dtype=acc_dtype), tCtO.layout) + tCtO = cute.make_tensor( + cute.recast_ptr(tmem_ptr + tmem_offset, dtype=acc_dtype), tCtO.layout + ) tmem_offset += tcgen05.find_tmem_tensor_col_offset(tCtO) - print(f"\t{tmem_offset} tmem cols used, {self.tmem_alloc_cols} tmem cols allocated") + print( + f"\t{tmem_offset} tmem cols used, {self.tmem_alloc_cols} tmem cols allocated" + ) assert tmem_offset <= self.tmem_alloc_cols ############################## # Exit early ############################## if exit_early: - if do_cluster_reduction and tidx == 0: - waiting_split_in_cluster = 0 - cute.arch.mbarrier_arrive( - m_cluster_full_ptr, - waiting_split_in_cluster, - arrive_count=mma_tile_n - ) + noop = None # early return not supported ############################## # TMA KV Dispatch @@ -720,14 +815,26 @@ class MixedInputFusedMultiHeadAttentionDecode: cute.arch.setmaxregister_decrease(self.mma_tma_regs) # Apply block tiler and slice - gK = cute.local_tile(mK, tiler=(blk_tile_s, blk_tile_d), coord=(None, 0, coord_hb)) # (TILE_S, TILE_D, #TILE_S) - gV = cute.local_tile(mV, tiler=(blk_tile_d, blk_tile_s), coord=(0, None, coord_hb)) # (TILE_D, TILE_S, #TILE_S) + gK = cute.local_tile( + mK, tiler=(blk_tile_s, blk_tile_d), coord=(None, 0, coord_hb) + ) # (TILE_S, TILE_D, #TILE_S) + gV = cute.local_tile( + mV, tiler=(blk_tile_d, blk_tile_s), coord=(0, None, coord_hb) + ) # (TILE_D, TILE_S, #TILE_S) # Apply MMA tiler and MMA partition - gK_mma = cute.flat_divide(gK, (mma_tile_m, mma_tile_k)) # (MMA_TILE_M, MMA_TILE_K, #TILE_SM, #TILE_DK, #TILE_S) - gV_mma = cute.flat_divide(gV, (mma_tile_m, mma_tile_k)) # (MMA_TILE_M, MMA_TILE_K, #TILE_DM, #TILE_SK, #TILE_S) - tAgK = thrblk_mma_kq.partition_A(gK_mma) # (MMA, #MMA_M, #MMA_K, #TILE_SM, #TILE_DK, #TILE_S) - tAgV = thrblk_mma_vp.partition_A(gV_mma) # (MMA, #MMA_M, #MMA_K, #TILE_DM, #TILE_SK, #TILE_S) + gK_mma = cute.flat_divide( + gK, (mma_tile_m, mma_tile_k) + ) # (MMA_TILE_M, MMA_TILE_K, #TILE_SM, #TILE_DK, #TILE_S) + gV_mma = cute.flat_divide( + gV, (mma_tile_m, mma_tile_k) + ) # (MMA_TILE_M, MMA_TILE_K, #TILE_DM, #TILE_SK, #TILE_S) + tAgK = thrblk_mma_kq.partition_A( + gK_mma + ) # (MMA, #MMA_M, #MMA_K, #TILE_SM, #TILE_DK, #TILE_S) + tAgV = thrblk_mma_vp.partition_A( + gV_mma + ) # (MMA, #MMA_M, #MMA_K, #TILE_DM, #TILE_SK, #TILE_S) # #TILE_SM=TILE_S/MMA_TILE_M, #TILE_HN=TILE_H/MMA_TILE_N, #TILE_DK=TILE_D/MMA_TILE_K # #TILE_DM=TILE_D/MMA_TILE_M, #TILE_HN=TILE_H/MMA_TILE_N, #TILE_SK=TILE_S/MMA_TILE_K @@ -739,14 +846,20 @@ class MixedInputFusedMultiHeadAttentionDecode: # TMA partition # (MMA, #MMA_M, #MMA_K, Rest...) -> (TMA, Rest...) tGSsK, tGSgK = cute.nvgpu.cpasync.tma_partition( - tma_atom_k, mcast_coord, mcast_layout, + tma_atom_k, + mcast_coord, + mcast_layout, smem_tensor=cute.group_modes(tAsK, 0, 3), - gmem_tensor=cute.group_modes(tAgK, 0, 3)) + gmem_tensor=cute.group_modes(tAgK, 0, 3), + ) tGSsV, tGSgV = cute.nvgpu.cpasync.tma_partition( - tma_atom_v, mcast_coord, mcast_layout, + tma_atom_v, + mcast_coord, + mcast_layout, smem_tensor=cute.group_modes(tAsV, 0, 3), - gmem_tensor=cute.group_modes(tAgV, 0, 3)) + gmem_tensor=cute.group_modes(tAgV, 0, 3), + ) # # Sequence loop @@ -762,11 +875,12 @@ class MixedInputFusedMultiHeadAttentionDecode: tma_atom_k, tGSgK_s[None, 0, dk], tGSsK[None, k_handle.index], - tma_bar_ptr=k_handle.barrier) + tma_bar_ptr=k_handle.barrier, + ) # Load V if s >= prefetch_tiles: - tGSgV_s = tGSgV[None, None, None, s-prefetch_tiles] + tGSgV_s = tGSgV[None, None, None, s - prefetch_tiles] for sk in cutlass.range_constexpr(tiles_sk): for dm in cutlass.range_constexpr(tiles_dm): v_handle = kv_producer.acquire_and_advance() @@ -774,7 +888,8 @@ class MixedInputFusedMultiHeadAttentionDecode: tma_atom_v, tGSgV_s[None, dm, sk], tGSsV[None, v_handle.index], - tma_bar_ptr=v_handle.barrier) + tma_bar_ptr=v_handle.barrier, + ) ############################## # TMA QO Dispatch @@ -785,39 +900,63 @@ class MixedInputFusedMultiHeadAttentionDecode: cute.arch.setmaxregister_decrease(self.mma_tma_regs) # Apply block tiler and slice - gQ = cute.local_tile(mQ, tiler=(blk_tile_h, blk_tile_d), coord=(coord_hr, 0, coord_hb)) # (TILE_H, TILE_D) - gO = cute.local_tile(mO, tiler=(blk_tile_d, blk_tile_h), coord=(0, coord_hr, coord_hb, kv_cluster_idx)) # (TILE_D, TILE_H) - gBS = cute.local_tile(mBS, tiler=(self.scaledim, blk_tile_s), coord=(0, None, coord_hb)) # (SCALE, TILE_S, #TILE_S) + gQ = cute.local_tile( + mQ, tiler=(blk_tile_h, blk_tile_d), coord=(coord_hr, 0, coord_hb) + ) # (TILE_H, TILE_D) + gO = cute.local_tile( + mO, + tiler=(blk_tile_d, blk_tile_h), + coord=(0, coord_hr, coord_hb, kv_split_idx), + ) # (TILE_D, TILE_H) + gBS = cute.local_tile( + mBS, tiler=(self.scaledim, blk_tile_s), coord=(0, None, coord_hb) + ) # (SCALE, TILE_S, #TILE_S) # Apply MMA tiler and MMA partition - gQ_mma = cute.flat_divide(gQ, (mma_tile_n, mma_tile_k)) # (MMA_TILE_N, MMA_TILE_K, #TILE_HN, #TILE_DK) - gO_mma = cute.flat_divide(gO, (mma_tile_m, mma_tile_n)) # (MMA_TILE_M, MMA_TILE_N, #TILE_DM, #TILE_HN) - tBgQ = thrblk_mma_kq.partition_B(gQ_mma) # (MMA, #MMA_N, #MMA_K, #TILE_HN, #TILE_DK) + gQ_mma = cute.flat_divide( + gQ, (mma_tile_n, mma_tile_k) + ) # (MMA_TILE_N, MMA_TILE_K, #TILE_HN, #TILE_DK) + gO_mma = cute.flat_divide( + gO, (mma_tile_m, mma_tile_n) + ) # (MMA_TILE_M, MMA_TILE_N, #TILE_DM, #TILE_HN) + tBgQ = thrblk_mma_kq.partition_B( + gQ_mma + ) # (MMA, #MMA_N, #MMA_K, #TILE_HN, #TILE_DK) # TMA partition tGSsQ, tGSgQ = cute.nvgpu.cpasync.tma_partition( - tma_atom_q, mcast_coord, mcast_layout, + tma_atom_q, + mcast_coord, + mcast_layout, smem_tensor=cute.group_modes(tBsQ, 0, 3), - gmem_tensor=cute.group_modes(tBgQ, 0, 3)) + gmem_tensor=cute.group_modes(tBgQ, 0, 3), + ) tSGsO, tSGgO = cute.nvgpu.cpasync.tma_partition( - tma_atom_o, mcast_coord, mcast_layout, + tma_atom_o, + mcast_coord, + mcast_layout, smem_tensor=cute.group_modes(sO_mma, 0, 2), - gmem_tensor=cute.group_modes(gO_mma, 0, 2)) + gmem_tensor=cute.group_modes(gO_mma, 0, 2), + ) tGSsBS, tGSgBS = cute.nvgpu.cpasync.tma_partition( - tma_atom_ks, mcast_coord, mcast_layout, + tma_atom_ks, + mcast_coord, + mcast_layout, smem_tensor=cute.group_modes(sBS, 0, 2), - gmem_tensor=cute.group_modes(gBS, 0, 2)) + gmem_tensor=cute.group_modes(gBS, 0, 2), + ) # Load Q for dk in cutlass.range_constexpr(tiles_dk): q_handle = q_producer.acquire_and_advance() cute.copy( tma_atom_q, - tGSgQ[None, 0, dk], # stages_q == tiles_dk by construction + tGSgQ[None, 0, dk], # stages_q == tiles_dk by construction tGSsQ[None, dk], - tma_bar_ptr=q_handle.barrier) + tma_bar_ptr=q_handle.barrier, + ) # Sequence Loop prefetch_tiles = prefetch_iters * kv_splits @@ -828,23 +967,22 @@ class MixedInputFusedMultiHeadAttentionDecode: tma_atom_ks, tGSgBS[None, s], tGSsBS[None, bs_handle.index], - tma_bar_ptr=bs_handle.barrier) + tma_bar_ptr=bs_handle.barrier, + ) if s >= prefetch_tiles: bs_handle = bs_producer.acquire_and_advance() cute.copy( tma_atom_vs, - tGSgBS[None, s-prefetch_tiles], + tGSgBS[None, s - prefetch_tiles], tGSsBS[None, bs_handle.index], - tma_bar_ptr=bs_handle.barrier) - + tma_bar_ptr=bs_handle.barrier, + ) + # Store O for dm in cutlass.range_constexpr(tiles_dm): - q_producer.acquire_and_advance() # Reuse Q load barriers to sync O store - cute.copy( - tma_atom_o, - tSGsO[None, dm, 0], - tSGgO[None, dm, 0]) + q_producer.acquire_and_advance() # Reuse Q load barriers to sync O store + cute.copy(tma_atom_o, tSGsO[None, dm, 0], tSGgO[None, dm, 0]) ############################## # Convert Dispatch @@ -854,65 +992,113 @@ class MixedInputFusedMultiHeadAttentionDecode: if cutlass.const_expr(self.use_reg_reconfig): cute.arch.setmaxregister_decrease(self.cvt_regs) - # Initialize for dual convert if necessary - convert_warpgroups = 1 + # Intermediate convert type + cvt_type = Float32 + if cutlass.const_expr(mma_dtype is cutlass.BFloat16 and + k_dtype in (cutlass.Int4, cutlass.Int8)): + cvt_type = mma_dtype + + # Initialize for multiple warpgroups if necessary convert_phase = 0 - if cutlass.const_expr(self.dual_convert): - assert tiles_dk % 2 == 0 - assert (tiles_dm * tiles_sk) % 2 == 0 - convert_warpgroups = 2 - convert_phase = warpgroup_idx % convert_warpgroups - if convert_phase == 1: + if cutlass.const_expr(self.convert_warpgroups > 1): + assert tiles_dk % self.convert_warpgroups == 0 + assert (tiles_dm * tiles_sk) % self.convert_warpgroups == 0 + convert_phase = warpgroup_idx % self.convert_warpgroups + for _ in cutlass.range(convert_phase): kv_consumer.advance() cvt_producer.advance() # Construct tiled copy and partition K - tmem_op_width = 32 - tmem_op_repeat = tcgen05.Repetition(mma_tile_k * mma_dtype.width // tmem_op_width) - tmem_store_atom_k = cute.make_copy_atom(tcgen05.St32x32bOp(tmem_op_repeat), mma_dtype) - tmem_store_k = tcgen05.make_tmem_copy(tmem_store_atom_k, tAtK_cvt[*mma_dice, 0]) - thr_store_k = tmem_store_k.get_slice(warpgroup_tidx) # tmem copy is always 128 threads + mma_k_bits = mma_tile_k * mma_dtype.width + tmem_store_atom_k = cute.make_copy_atom( + tcgen05.St16x256bOp(tcgen05.Repetition(mma_k_bits // 256)), + mma_dtype, + ) + smem_load_atom_k = cute.make_copy_atom( + cute.nvgpu.warp.LdMatrix8x16x8bOp( + num_matrices=4, + unpack_bits=(k_dtype.width if k_dtype.width < 8 else None)), + kv_smem_dtype, + ) - tKsK = thr_store_k.partition_S(tAsK) + tmem_store_k = tcgen05.make_tmem_copy( + tmem_store_atom_k, tAtK_cvt[mma_dice+(0,)] + ) + thr_store_k = tmem_store_k.get_slice(warpgroup_tidx) + tKrK_cvt_shape = thr_store_k.partition_S(tAtK_cvt).shape[:-1] tKtK_cvt = thr_store_k.partition_D(tAtK_cvt) - # Construct tiled copy and partition V - tmem_op_width = 128 - tmem_op_repeat = tcgen05.Repetition(mma_tile_k * mma_dtype.width // tmem_op_width) - tmem_store_atom_v = cute.make_copy_atom(tcgen05.St16x128bOp(tmem_op_repeat), mma_dtype) - tmem_store_v = tcgen05.make_tmem_copy(tmem_store_atom_v, tAtV_cvt[*mma_dice, 0]) - thr_store_v = tmem_store_v.get_slice(warpgroup_tidx) + smem_load_k = cute.make_tiled_copy_S(smem_load_atom_k, tmem_store_k) + thr_load_k = smem_load_k.get_slice(warpgroup_tidx) + tKsK = thr_load_k.partition_S(tAsK) + tKrK_shape = thr_load_k.partition_D(tAsK).shape[:-1] + + # Construct tiled copy and partition V + tmem_store_atom_v = cute.make_copy_atom( + tcgen05.St16x256bOp(tcgen05.Repetition(mma_k_bits // 256)), mma_dtype + ) + smem_load_op_v = cute.nvgpu.warp.LdMatrix16x16x8bOp( + transpose=True, + num_matrices=2, + unpack_bits=(v_dtype.width if v_dtype.width < 8 else None), + ) + smem_load_atom_v = cute.make_copy_atom(smem_load_op_v, kv_smem_dtype) + + tmem_store_v = tcgen05.make_tmem_copy( + tmem_store_atom_v, tAtV_cvt[mma_dice+(0,)] + ) + thr_store_v = tmem_store_v.get_slice(warpgroup_tidx) + tVrV_cvt_shape = thr_store_v.partition_S(tAtV_cvt).shape[:-1] + tVtV_cvt = thr_store_v.partition_D(tAtV_cvt) - smem_load_atom_v = cute.make_copy_atom(cute.nvgpu.warp.LdMatrix16x16x8bOp(num_matrices=2), v_dtype) smem_load_v = cute.make_tiled_copy_S(smem_load_atom_v, tmem_store_v) thr_load_v = smem_load_v.get_slice(warpgroup_tidx) - tVsV = thr_load_v.partition_S(tAsV) - tVrV_shape = smem_load_v.get_slice(0).partition_D(tAsV).shape[:-1] - tVrV_cvt_shape = thr_store_v.get_slice(0).partition_S(tAtV_cvt).shape[:-1] - tVtV_cvt = thr_store_v.partition_D(tAtV_cvt) + tVrV_shape = thr_load_v.partition_D(tAsV).shape[:-1] # Partition KS - K block scale sKS_layout = cute.make_layout( - shape=(blk_tile_s, (self.block_scaledim, self.scaledim), self.bs_stages), - stride=(self.scaledim, (0, 1), blk_tile_s * self.scaledim)) - sKS = cute.make_tensor(sBS.iterator, sKS_layout) # (TILE_S, TILE_D, bs_stages) - # (MMA_TILE_M, MMA_TILE_K, (#TILE_SM, #TILE_DK), bs_stages) - sKS_mma = cute.group_modes(cute.flat_divide(sKS, (mma_tile_m, mma_tile_k)), 2, 4) - tAsKS = thrblk_mma_kq.partition_A(sKS_mma) # (MMA, #MMA_M, #MMA_K, (#TILE_SM, #TILE_DK), bs_stages) - tKsKS = thr_store_k.partition_S(tAsKS) # (CPY, CPY_MMA, CPY_M, CPY_K, #TILE, bs_stages) - tKrKS = cute.make_rmem_tensor_like(tKsKS[*cpy_dice, None, 0]) # (CPY, CPY_MMA, CPY_M, CPY_K, #TILE) + shape=( + blk_tile_s, + (self.block_scaledim, self.scaledim), + self.bs_stages, + ), + stride=(self.scaledim, (0, 1), blk_tile_s * self.scaledim), + ) + sKS = cute.make_tensor( + sBS.iterator, sKS_layout + ) # (TILE_S, TILE_D, bs_stages) + sKS_mma = cute.group_modes( + cute.flat_divide(sKS, (mma_tile_m, mma_tile_k)), 2, 4 + ) # (MMA_TILE_M, MMA_TILE_K, (#TILE_SM, #TILE_DK), bs_stages) + tAsKS = thrblk_mma_kq.partition_A( + sKS_mma + ) # (MMA, #MMA_M, #MMA_K, (#TILE_SM, #TILE_DK), bs_stages) + tKsKS = thr_load_k.partition_D( + tAsKS + ) # (CPY, CPY_MMA, CPY_M, CPY_K, #TILE, bs_stages) # Partition VS - V block scale sVS_layout = cute.make_layout( - shape=((self.block_scaledim, self.scaledim), blk_tile_s, self.bs_stages), - stride=((0, 1), self.scaledim, blk_tile_s * self.scaledim)) - sVS = cute.make_tensor(sBS.iterator, sVS_layout) # (TILE_D, TILE_S, bs_stages) - # (MMA_TILE_M, MMA_TILE_K, (#TILE_DM, #TILE_SK), bs_stages) - sVS_mma = cute.group_modes(cute.flat_divide(sVS, (mma_tile_m, mma_tile_k)), 2, 4) - tAsVS = thrblk_mma_vp.partition_A(sVS_mma) # (MMA, #MMA_M, #MMA_K, (#TILE_DM, #TILE_SK), bs_stages) - tVsVS = thr_load_v.partition_D(tAsVS) # (CPY, CPY_MMA, CPY_M, CPY_K, #TILE, bs_stages) - tVrVS = cute.make_rmem_tensor_like(tVsVS[*cpy_dice, None, 0]) # (CPY, CPY_MMA, CPY_M, CPY_K, #TILE) + shape=( + (self.block_scaledim, self.scaledim), + blk_tile_s, + self.bs_stages, + ), + stride=((0, 1), self.scaledim, blk_tile_s * self.scaledim), + ) + sVS = cute.make_tensor( + sBS.iterator, sVS_layout + ) # (TILE_D, TILE_S, bs_stages) + sVS_mma = cute.group_modes( + cute.flat_divide(sVS, (mma_tile_m, mma_tile_k)), 2, 4 + ) # (MMA_TILE_M, MMA_TILE_K, (#TILE_DM, #TILE_SK), bs_stages) + tAsVS = thrblk_mma_vp.partition_A( + sVS_mma + ) # (MMA, #MMA_M, #MMA_K, (#TILE_DM, #TILE_SK), bs_stages) + tVsVS = thr_load_v.partition_D( + tAsVS + ) # (CPY, CPY_MMA, CPY_M, CPY_K, #TILE, bs_stages) # # Sequence loop @@ -921,62 +1107,94 @@ class MixedInputFusedMultiHeadAttentionDecode: if s < iters_s: # Load K scale bs_handle = bs_consumer.wait_and_advance() - cute.autovec_copy(tKsKS[*cpy_dice, None, bs_handle.index], tKrKS) + tKrKS = cute.make_rmem_tensor_like( + tKsKS[cpy_dice + (None, 0)] + ) # 'like' preserves 0 strides + cute.autovec_copy(tKsKS[cpy_dice + (None, bs_handle.index)], tKrKS) cute.arch.fence_view_async_shared() bs_handle.release() # Convert and scale K - for dk in cutlass.range(tiles_dk // convert_warpgroups, unroll=2): - tKrK = cute.make_rmem_tensor(tKsK.shape[:-1], k_dtype) - tKrK_cvt = cute.make_rmem_tensor(tKsK.shape[:-1], mma_dtype) + for dk in cutlass.range(tiles_dk // self.convert_warpgroups, unroll=2): + tKrK = cute.make_rmem_tensor(tKrK_shape, kv_smem_dtype) + tKrK_cvt = cute.make_rmem_tensor(tKrK_cvt_shape, mma_dtype) kv_handle = kv_consumer.wait_and_advance() - cute.autovec_copy(tKsK[*cpy_dice, kv_handle.index], tKrK) + cute.copy(thr_load_k, tKsK[cpy_dice + (kv_handle.index,)], tKrK) cute.arch.fence_view_async_shared() kv_handle.release() - coord_dk = dk * convert_warpgroups + convert_phase - scale_k = tKrKS[*cpy_dice, coord_dk].load() - tKrK_cvt.store(tKrK.load().to(mma_dtype) * scale_k) + # Sign extend unpacked int4 to int8 + if cutlass.const_expr(k_dtype is cutlass.Int4): + tKrK_unpacked_i4_vec = tKrK.load().maybe_downcast() + tKrK_i8_vec = cute.arch.sext_unpacked_i4_i8_intrinsic( + tKrK_unpacked_i4_vec, cute.size(tKrK_shape) + ) + tKrK.store( + cute.TensorSSA(tKrK_i8_vec, tKrK_shape, cutlass.Int8) + ) + + coord_dk = dk * self.convert_warpgroups + convert_phase + scale_k = tKrKS[cpy_dice + (coord_dk,)].load() + tKrK_ssa = tKrK.load().to(cvt_type).to(mma_dtype) * scale_k + tKrK_cvt.store(tKrK_ssa.reshape(tKrK_cvt_shape)) cvt_handle = cvt_producer.acquire_and_advance() - cute.copy(thr_store_k, tKrK_cvt, tKtK_cvt[*cpy_dice, cvt_handle.index]) + cute.copy( + thr_store_k, tKrK_cvt, tKtK_cvt[cpy_dice + (cvt_handle.index,)] + ) cute.arch.fence_view_async_tmem_store() cvt_handle.commit() - # Advance again for dual warpgroups - if cutlass.const_expr(self.dual_convert): + # Advance again for multiple warpgroups + for _ in cutlass.range_constexpr(self.convert_warpgroups - 1): kv_consumer.advance() cvt_producer.advance() if s >= prefetch_iters: # Load V scale bs_handle = bs_consumer.wait_and_advance() - cute.autovec_copy(tVsVS[*cpy_dice, None, bs_handle.index], tVrVS) + tVrVS = cute.make_rmem_tensor_like( + tVsVS[cpy_dice + (None, 0)] + ) # 'like' preserves 0 strides + cute.autovec_copy(tVsVS[cpy_dice + (None, bs_handle.index)], tVrVS) cute.arch.fence_view_async_shared() bs_handle.release() # Convert and scale V - for dmsk in cutlass.range(tiles_dm * tiles_sk // convert_warpgroups, unroll=2): - tVrV = cute.make_rmem_tensor(tVrV_shape, v_dtype) + for dmsk in cutlass.range(tiles_dm * tiles_sk // self.convert_warpgroups, unroll=2): + tVrV = cute.make_rmem_tensor(tVrV_shape, kv_smem_dtype) tVrV_cvt = cute.make_rmem_tensor(tVrV_cvt_shape, mma_dtype) kv_handle = kv_consumer.wait_and_advance() - cute.copy(thr_load_v, tVsV[*cpy_dice, kv_handle.index], tVrV) + cute.copy(thr_load_v, tVsV[cpy_dice + (kv_handle.index,)], tVrV) cute.arch.fence_view_async_shared() kv_handle.release() - coord_dmsk = dmsk * convert_warpgroups + convert_phase - scale_v = tVrVS[*cpy_dice, coord_dmsk].load() - tVrV_cvt.store(tVrV.load().to(mma_dtype) * scale_v) + # Sign extend unpacked int4 to int8 + if cutlass.const_expr(v_dtype is cutlass.Int4): + tVrV_unpacked_i4_vec = tVrV.load().maybe_downcast() + tVrV_i8_vec = cute.arch.sext_unpacked_i4_i8_intrinsic( + tVrV_unpacked_i4_vec, cute.size(tVrV_shape) + ) + tVrV.store( + cute.TensorSSA(tVrV_i8_vec, tVrV_shape, cutlass.Int8) + ) + + coord_dmsk = dmsk * self.convert_warpgroups + convert_phase + scale_v = tVrVS[cpy_dice + (coord_dmsk,)].load() + tVrV_ssa = tVrV.load().to(cvt_type).to(mma_dtype) * scale_v + tVrV_cvt.store(tVrV_ssa.reshape(tVrV_cvt_shape)) cvt_handle = cvt_producer.acquire_and_advance() - cute.copy(thr_store_v, tVrV_cvt, tVtV_cvt[*cpy_dice, cvt_handle.index]) + cute.copy( + thr_store_v, tVrV_cvt, tVtV_cvt[cpy_dice + (cvt_handle.index,)] + ) cute.arch.fence_view_async_tmem_store() cvt_handle.commit() - # Advance again for dual warpgroups - if cutlass.const_expr(self.dual_convert): + # Advance again for multiple warpgroups + for _ in cutlass.range_constexpr(self.convert_warpgroups - 1): kv_consumer.advance() cvt_producer.advance() @@ -996,7 +1214,7 @@ class MixedInputFusedMultiHeadAttentionDecode: q_consumer.wait_and_advance() # Sequence loop - s_token = True # Producer always acquires first + s_token = True # Producer always acquires first for s in cutlass.range(iters_s): # BMM1 k_token = cvt_consumer.try_wait() @@ -1007,14 +1225,16 @@ class MixedInputFusedMultiHeadAttentionDecode: k_handle = cvt_consumer.wait_and_advance(k_token) # Signal BMM2 to start if is_last_iter: - cute.arch.barrier_arrive(barrier_id=mma_kq_nbar_id, number_of_threads=64) + cute.arch.barrier_arrive( + barrier_id=mma_kq_nbar_id, number_of_threads=64 + ) for mma_k in cutlass.range_constexpr(tAtK_cvt.shape[2]): cute.gemm( tiled_mma_kq, - tCtS[*mma_dice, s_handle.index], + tCtS[mma_dice + (s_handle.index,)], tAtK_cvt[None, None, mma_k, k_handle.index], tBsQ_desc[None, None, mma_k, dk], - tCtS[*mma_dice, s_handle.index], + tCtS[mma_dice + (s_handle.index,)], ) if dk == 0 and mma_k == 0: tiled_mma_kq.set(tcgen05.Field.ACCUMULATE, True) @@ -1030,7 +1250,6 @@ class MixedInputFusedMultiHeadAttentionDecode: cute.arch.barrier(barrier_id=mma_vp_nbar_id, number_of_threads=64) s_token = s_producer.try_acquire() - ############################## # MMA VP Dispatch ############################## @@ -1043,13 +1262,14 @@ class MixedInputFusedMultiHeadAttentionDecode: tiled_mma_vp.set(tcgen05.Field.ACCUMULATE, True) tBsP_desc = thrblk_mma_vp.make_fragment_B(tBsP_nk) + # Advance and wait for BMM1 for _ in cutlass.range_constexpr(tiles_dk): cvt_consumer.advance() cute.arch.barrier(barrier_id=mma_kq_nbar_id, number_of_threads=64) # Sequence loop p_token = False - o_token = True # Producer always acquires first + o_token = True # Producer always acquires first for s in cutlass.range(iters_s): # Advance and wait for BMM1 if s < iters_s - 1: @@ -1068,14 +1288,16 @@ class MixedInputFusedMultiHeadAttentionDecode: v_handle = cvt_consumer.wait_and_advance(v_token) # Signal BMM1 to start if is_last_iter: - cute.arch.barrier_arrive(barrier_id=mma_vp_nbar_id, number_of_threads=64) + cute.arch.barrier_arrive( + barrier_id=mma_vp_nbar_id, number_of_threads=64 + ) for mma_k in cutlass.range_constexpr(tAtV_cvt.shape[2]): cute.gemm( tiled_mma_vp, - tCtO[*mma_dice, dm, 0], + tCtO[mma_dice + (dm, 0)], tAtV_cvt[None, None, mma_k, v_handle.index], tBsP_desc[None, None, mma_k, sk, p_handle.index], - tCtO[*mma_dice, dm, 0], + tCtO[mma_dice + (dm, 0)], ) v_handle.release() if not is_last_iter: @@ -1099,43 +1321,70 @@ class MixedInputFusedMultiHeadAttentionDecode: # Construct tiled copies tmem_op_width = 32 - tmem_op_repeat = tcgen05.Repetition(mma_tile_n * acc_dtype.width // tmem_op_width) - tmem_load_atom_s = cute.make_copy_atom(tcgen05.Ld32x32bOp(tmem_op_repeat), acc_dtype) - tmem_load_s = tcgen05.make_tmem_copy(tmem_load_atom_s, tCtS[*mma_dice, 0]) + tmem_op_repeat = tcgen05.Repetition( + mma_tile_n * acc_dtype.width // tmem_op_width + ) + tmem_load_atom_s = cute.make_copy_atom( + tcgen05.Ld32x32bOp(tmem_op_repeat), acc_dtype + ) + tmem_load_s = tcgen05.make_tmem_copy(tmem_load_atom_s, tCtS[mma_dice + (0,)]) thr_load_s = tmem_load_s.get_slice(warpgroup_tidx) - tmem_store_atom_o = cute.make_copy_atom(tcgen05.St32x32bOp(tmem_op_repeat), o_dtype) - tmem_store_o = tcgen05.make_tmem_copy(tmem_store_atom_o, tCtO[*mma_dice, 0, 0]) + tmem_store_atom_o = cute.make_copy_atom( + tcgen05.St32x32bOp(tmem_op_repeat), o_dtype + ) + tmem_store_o = tcgen05.make_tmem_copy( + tmem_store_atom_o, tCtO[mma_dice + (0, 0)] + ) thr_store_o = tmem_store_o.get_slice(warpgroup_tidx) # Partition S and P - tStS = thr_load_s.partition_S(tCtS) # (CPY, #CPY_MMA, #CPY_M, #CPY_N, stages_sp) - tSsP = thr_load_s.partition_D(tCsP) # (CPY, #CPY_MMA, #CPY_M, #CPY_N, stages_sp) + tStS = thr_load_s.partition_S( + tCtS + ) # (CPY, #CPY_MMA, #CPY_M, #CPY_N, stages_sp) + tSsP = thr_load_s.partition_D( + tCsP + ) # (CPY, #CPY_MMA, #CPY_M, #CPY_N, stages_sp) # Partition O - tStO = thr_load_s.partition_S(tCtO) # (CPY, #CPY_MMA, #CPY_M, #CPY_N, #TILE_DM, #TILE_HN) - tSsO = thr_load_s.partition_D(tCsO) # (CPY, #CPY_MMA, #CPY_M, #CPY_N, #TILE_DM, #TILE_HN) + tStO = thr_load_s.partition_S( + tCtO + ) # (CPY, #CPY_MMA, #CPY_M, #CPY_N, #TILE_DM, #TILE_HN) + tSsO = thr_load_s.partition_D( + tCsO + ) # (CPY, #CPY_MMA, #CPY_M, #CPY_N, #TILE_DM, #TILE_HN) tSrO = cute.make_rmem_tensor_like(tSsO) # Partition colmax and initialize in RF - tSsM = thr_load_s.partition_D(tCsM) # (CPY, #CPY_MMA, #CPY_M, #CPY_N) - tSsM_cluster = thr_load_s.partition_D(tCsM_cluster) + tSsM = thr_load_s.partition_D(tCsM) # (CPY, #CPY_MMA, #CPY_M, #CPY_N) tSrM_prev = cute.make_rmem_tensor_like(tSsM) tSrM_prev.fill(-Float32.inf) # Partition colsum and initialize in RF # Each thread maintains a local colsum in RF, smem reduction happens after loop - tSsL = thr_load_s.partition_D(tCsL) # (CPY, #CPY_MMA, #CPY_M, #CPY_N, WARPS) - tSrL = cute.make_rmem_tensor_like(tSsL[*cpy_dice, 0]) + tSsL = thr_load_s.partition_D( + tCsL + ) # (CPY, #CPY_MMA, #CPY_M, #CPY_N, WARPS) + tSrL = cute.make_rmem_tensor_like(tSsL[cpy_dice + (0,)]) tSrL.fill(Float32(0)) assert warp_threads >= cute.size(tSsM) # get gmem colmax + colsum to store to inbound_hr = coord_hr * blk_tile_h + lane_idx < mM.shape[0] - gM = cute.local_tile(mM, tiler=(mma_tile_n, 1), coord=(coord_hr, coord_hb)) # (TILE_H) = (MMA_TILE_N) - gM_partial = cute.local_tile(mM_partial, tiler=(mma_tile_n, 1), coord=(coord_hr, coord_hb, kv_cluster_idx)) - gL_partial = cute.local_tile(mL_partial, tiler=(mma_tile_n, 1), coord=(coord_hr, coord_hb, kv_cluster_idx)) + gM = cute.local_tile( + mM, tiler=(mma_tile_n, 1), coord=(coord_hr, coord_hb) + ) # (TILE_H) = (MMA_TILE_N) + gM_partial = cute.local_tile( + mM_partial, + tiler=(mma_tile_n, 1), + coord=(coord_hr, coord_hb, kv_split_idx), + ) + gL_partial = cute.local_tile( + mL_partial, + tiler=(mma_tile_n, 1), + coord=(coord_hr, coord_hb, kv_split_idx), + ) # Initialize O tSrO.fill(Float32(0)) @@ -1146,7 +1395,9 @@ class MixedInputFusedMultiHeadAttentionDecode: tSsM[lane_idx] = -Float32.inf if warpgroup_widx == 1 and lane_store_max: tSsL[lane_idx] = Float32(0) - cute.arch.barrier(barrier_id=softmax_nbar_id, number_of_threads=warpgroup_threads) + cute.arch.barrier( + barrier_id=softmax_nbar_id, number_of_threads=warpgroup_threads + ) # # Sequence loop @@ -1155,15 +1406,15 @@ class MixedInputFusedMultiHeadAttentionDecode: # Load S from tmem s_handle = s_consumer.wait_and_advance() tSrS = cute.make_rmem_tensor(tSsP.shape[:-1], acc_dtype) - cute.copy(tmem_load_s, tStS[*cpy_dice, s_handle.index], tSrS) + cute.copy(tmem_load_s, tStS[cpy_dice + (s_handle.index,)], tSrS) cute.arch.fence_view_async_tmem_load() s_handle.release() # Reduce colmax in warp RF tSrM = cute.make_rmem_tensor_like(tSsM) - tSrM_lane = Float32(0) # Avoid dynamic register indexing + tSrM_lane = Float32(0) # Avoid dynamic register indexing for i in cutlass.range_constexpr(cute.size(tSrS)): - tSrM[i] = warp_fmax(tSrS[i]) + tSrM[i] = cute.arch.warp_redux_sync(tSrS[i], kind="fmax", nan=True) if i == lane_idx: tSrM_lane = tSrM[i] @@ -1172,44 +1423,56 @@ class MixedInputFusedMultiHeadAttentionDecode: self.smem_fmax(tSsM.iterator + tSsM.layout(lane_idx), tSrM_lane) # Wait for colmax then load - cute.arch.barrier(barrier_id=softmax_nbar_id, number_of_threads=warpgroup_threads) + cute.arch.barrier( + barrier_id=softmax_nbar_id, number_of_threads=warpgroup_threads + ) cute.autovec_copy(tSsM, tSrM) # Compute online softmax tSrP = cute.make_rmem_tensor(tSsP.shape[:-1], mma_dtype) if cutlass.const_expr(use_tensor_ssa_math): tSrP_f32 = exp2(scale_qs_log2_e * (tSrS.load() - tSrM.load())) - tSrP.store(tSrP_f32.to(mma_dtype)) # convert + tSrP.store(tSrP_f32.to(mma_dtype)) # convert else: tSrP_f32 = cute.make_rmem_tensor(tSrS.shape, acc_dtype) for i in cutlass.range_constexpr(0, cute.size(tSrS), 2): - p_f32x2 = fadd2((tSrS[i], tSrS[i+1]), (-tSrM[i], -tSrM[i+1])) + p_f32x2 = fadd2( + (tSrS[i], tSrS[i + 1]), (-tSrM[i], -tSrM[i + 1]) + ) p_f32x2 = fmul2(p_f32x2, (scale_qs_log2_e, scale_qs_log2_e)) tSrP_f32[i] = exp2(p_f32x2[0]) - tSrP_f32[i+1] = exp2(p_f32x2[1]) + tSrP_f32[i + 1] = exp2(p_f32x2[1]) tSrP.store(tSrP_f32.load().to(mma_dtype)) # Store P to smem p_handle = p_producer.acquire_and_advance() - cute.autovec_copy(tSrP, tSsP[*cpy_dice, p_handle.index]) + cute.autovec_copy(tSrP, tSsP[cpy_dice + (p_handle.index,)]) cute.arch.fence_view_async_shared() p_handle.commit() # Compute correction and correct colsum if cutlass.const_expr(use_tensor_ssa_math): - correction = exp2(scale_qs_log2_e * (tSrM_prev.load() - tSrM.load())) + correction = exp2( + scale_qs_log2_e * (tSrM_prev.load() - tSrM.load()) + ) tSrL.store(tSrL.load() * correction + tSrP_f32) else: correction = cute.make_rmem_tensor_like(tSrM) for i in cutlass.range_constexpr(0, cute.size(tSrM), 2): - c_f32x2 = fadd2((tSrM_prev[i], tSrM_prev[i+1]), (-tSrM[i], -tSrM[i+1])) + c_f32x2 = fadd2( + (tSrM_prev[i], tSrM_prev[i + 1]), (-tSrM[i], -tSrM[i + 1]) + ) c_f32x2 = fmul2(c_f32x2, (scale_qs_log2_e, scale_qs_log2_e)) c_f32x2 = (exp2(c_f32x2[0]), exp2(c_f32x2[1])) correction[i] = c_f32x2[0] - correction[i+1] = c_f32x2[1] - l_f32x2 = ffma2(c_f32x2, (tSrL[i], tSrL[i+1]), (tSrP_f32[i], tSrP_f32[i+1])) + correction[i + 1] = c_f32x2[1] + l_f32x2 = ffma2( + c_f32x2, + (tSrL[i], tSrL[i + 1]), + (tSrP_f32[i], tSrP_f32[i + 1]), + ) tSrL[i] = l_f32x2[0] - tSrL[i+1] = l_f32x2[1] + tSrL[i + 1] = l_f32x2[1] # Correct O if s > 0: @@ -1218,15 +1481,18 @@ class MixedInputFusedMultiHeadAttentionDecode: # Apply correction for dm in cutlass.range_constexpr(tiles_dm): - tSrO_dm = cute.make_rmem_tensor_like(tSsO[*cpy_dice, 0, 0]) - cute.copy(thr_load_s, tStO[*cpy_dice, dm, 0], tSrO_dm) + tSrO_dm = cute.make_rmem_tensor_like(tSsO[cpy_dice + (0, 0)]) + cute.copy(thr_load_s, tStO[cpy_dice + (dm, 0)], tSrO_dm) for i in cutlass.range_constexpr(0, cute.size(tSrO_dm), 2): - o_f32x2 = fmul2((tSrO_dm[i], tSrO_dm[i+1]), (correction[i], correction[i+1])) + o_f32x2 = fmul2( + (tSrO_dm[i], tSrO_dm[i + 1]), + (correction[i], correction[i + 1]), + ) tSrO_dm[i] = o_f32x2[0] - tSrO_dm[i+1] = o_f32x2[1] + tSrO_dm[i + 1] = o_f32x2[1] - cute.copy(thr_store_o, tSrO_dm, tStO[*cpy_dice, dm, 0]) + cute.copy(thr_store_o, tSrO_dm, tStO[cpy_dice + (dm, 0)]) # Notify MMA cute.arch.fence_view_async_tmem_store() @@ -1248,117 +1514,38 @@ class MixedInputFusedMultiHeadAttentionDecode: # Store partial colsum in smem if lane_store_max: - tSsL[*cpy_dice, warpgroup_widx][lane_idx] = tSrL_lane + tSsL[cpy_dice + (warpgroup_widx,)][lane_idx] = tSrL_lane - # Reduce cluster colmax and correct O - if do_cluster_reduction: - # Reduce cluster colmax - if warpgroup_widx == 0: - if lane_store_max: - self.dsmem_fmax( - sM_cluster.iterator + sM_layout((0, lane_idx)), - sM[(0, lane_idx)], - m_cluster_full_ptr - ) + # Wait for colsum + cute.arch.barrier( + barrier_id=softmax_nbar_id, number_of_threads=warpgroup_threads + ) - # split 0 waits for cluster colmax to finish reduction - # other splits wait for split 0 to notify cluster colmax is ready - cute.arch.mbarrier_wait(m_cluster_full_ptr, phase=0) + if warpgroup_widx == 0 and lane_store_max and inbound_hr: + # Load colsum and colmax + sL_lane_wg = sL[0, lane_idx, None] + sL_lane = ( + sL_lane_wg[0] + sL_lane_wg[1] + sL_lane_wg[2] + sL_lane_wg[3] + ) + sM_lane = sM[0, lane_idx] - if warpgroup_widx == 0: - if kv_split_in_cluster == 0: - # notify other splits that cluster colmax is ready in split 0 smem - if lane_idx > 0 and lane_idx < kv_cluster_dim: - waiting_split_in_cluster = lane_idx - cute.arch.mbarrier_arrive( - m_cluster_full_ptr, waiting_split_in_cluster, arrive_count=1 - ) - else: - # other splits copy cluster colmax into local smem - if lane_store_max: - sM_cluster[0, lane_idx] = self.dsmem_load( - sM_cluster.iterator + sM_layout((0, lane_idx)) - ) + # Scale colmax + sM_lane = sM_lane * scale_qs - # warpgroup waits for cluster colmax to load into local smem - cute.arch.barrier(barrier_id=softmax_nbar_id, number_of_threads=warpgroup_threads) + # Store colsum and colmax + gL_partial[lane_idx] = sL_lane + gM_partial[lane_idx] = sM_lane + self.gmem_fmax(gM.iterator + gM.layout(lane_idx), sM_lane) - if warpgroup_widx == 0 and lane_store_max and inbound_hr: - # Load colsum and colmax - sL_lane_wg = sL[0, lane_idx, None] - sL_lane = sL_lane_wg[0] + sL_lane_wg[1] + sL_lane_wg[2] + sL_lane_wg[3] - sM_prev_lane = sM[0, lane_idx] - sM_lane = sM_cluster[0, lane_idx] - - # Correct colsum and scale colmax - correction = exp2(scale_qs_log2_e * (sM_prev_lane - sM_lane)) - sL_lane = sL_lane * correction - sM_lane = sM_lane * scale_qs - - # Store colsum and colmax - self.gmem_fadd( - gL_partial.iterator + gL_partial.layout(lane_idx), sL_lane - ) - self.gmem_fmax(gM.iterator + gM.layout(lane_idx), sM_lane) - if kv_split_in_cluster == 0: - gM_partial[lane_idx] = sM_lane - - # Load cluster colmax - tSrM = cute.make_rmem_tensor_like(tSsM) - cute.autovec_copy(tSsM_cluster, tSrM) - - # Wait and load O - o_handle = o_consumer.wait_and_advance() - cute.copy(thr_load_s, tStO, tSrO) - cute.arch.fence_view_async_tmem_load() - o_handle.release() # Final release signals tmem dealloc - - # Apply cluster correction - if cutlass.const_expr(use_tensor_ssa_math): - correction = exp2(scale_qs_log2_e * (tSrM_prev.load() - tSrM.load())) - tSrO.store(tSrO.load() * correction) - else: - correction = cute.make_rmem_tensor_like(tSsM) - for i in cutlass.range_constexpr(0, cute.size(tSrM), 2): - c_f32x2 = fadd2((tSrM_prev[i], tSrM_prev[i+1]), (-tSrM[i], -tSrM[i+1])) - c_f32x2 = fmul2(c_f32x2, (scale_qs_log2_e, scale_qs_log2_e)) - correction[i] = exp2(c_f32x2[0]) - correction[i+1] = exp2(c_f32x2[1]) - - for i in cutlass.range_constexpr(0, cute.size(tSrO), cute.size(correction)): - for j in cutlass.range_constexpr(0, cute.size(correction), 2): - o_f32x2 = fmul2((tSrO[i+j], tSrO[i+j+1]), (correction[j], correction[j+1])) - tSrO[i+j] = o_f32x2[0] - tSrO[i+j+1] = o_f32x2[1] - - # Wait and load O without cluster correction - else: - # Wait for colsum - cute.arch.barrier(barrier_id=softmax_nbar_id, number_of_threads=warpgroup_threads) - - if warpgroup_widx == 0 and lane_store_max and inbound_hr: - # Load colsum and colmax - sL_lane_wg = sL[0, lane_idx, None] - sL_lane = sL_lane_wg[0] + sL_lane_wg[1] + sL_lane_wg[2] + sL_lane_wg[3] - sM_lane = sM[0, lane_idx] - - # Scale colmax - sM_lane = sM_lane * scale_qs - - # Store colsum and colmax - gL_partial[lane_idx] = sL_lane - gM_partial[lane_idx] = sM_lane - self.gmem_fmax(gM.iterator + gM.layout(lane_idx), sM_lane) - - o_handle = o_consumer.wait_and_advance() - cute.copy(thr_load_s, tStO, tSrO) - cute.arch.fence_view_async_tmem_load() - o_handle.release() # Final release signals tmem dealloc + o_handle = o_consumer.wait_and_advance() + cute.copy(thr_load_s, tStO, tSrO) + cute.arch.fence_view_async_tmem_load() + o_handle.release() # Final release signals tmem dealloc # Store O to smem for dm in cutlass.range_constexpr(tiles_dm): - tOrO_dm = tSrO[*cpy_dice, dm, 0] - tOsO_dm = tSsO[*cpy_dice, dm, 0] + tOrO_dm = tSrO[cpy_dice + (dm, 0)] + tOsO_dm = tSsO[cpy_dice + (dm, 0)] cute.autovec_copy(tOrO_dm, tOsO_dm) cute.arch.fence_view_async_shared() @@ -1367,25 +1554,18 @@ class MixedInputFusedMultiHeadAttentionDecode: q_consumer.release() q_consumer.advance() - - # Ensure split 0 doesn't exit before all splits read cluster colmax - if do_cluster_reduction: - cute.arch.cluster_arrive_relaxed() - if kv_split_in_cluster == 0 and warp_idx == self.tma_qo_warp_id: - cute.arch.cluster_wait() - return @staticmethod @cute.kernel def reduction( - o : cute.Tensor, - m : cute.Tensor, - l : cute.Tensor, - o_partial : cute.Tensor, - m_partial : cute.Tensor, - l_partial : cute.Tensor, - scale_o : Float32, + o: cute.Tensor, + m: cute.Tensor, + l: cute.Tensor, + o_partial: cute.Tensor, + m_partial: cute.Tensor, + l_partial: cute.Tensor, + scale_o: Float32, ): d_blk_idx, coord_h, coord_b = cute.arch.block_idx() d_per_blk, _, _ = cute.arch.block_dim() @@ -1415,48 +1595,9 @@ class MixedInputFusedMultiHeadAttentionDecode: @staticmethod @cute.jit - def _mapa(ptr : Pointer, cta_rank_in_cluster : Int32 = 0): - llvm_ptr = ptr.llvm_ptr - return nvvm.mapa_shared_cluster( - llvm_ptr.type, - llvm_ptr, - Int32(cta_rank_in_cluster).ir_value(), - ) - - @cute.jit - def dsmem_load(self, val_ptr: Pointer): - val_llvm_ptr = self._mapa(val_ptr, 0) - - ret = llvm.inline_asm( - Float32.mlir_type, - [val_llvm_ptr], - "ld.relaxed.cta.shared::cluster.f32 $0, [$1];", - "=f,r", - has_side_effects=False, - is_align_stack=False, - asm_dialect=llvm.AsmDialect.AD_ATT, - ) - - return Float32(ret) - - @staticmethod - @cute.jit - def warp_fmax(val : Float32): - ret = llvm.inline_asm( - Float32.mlir_type, - [val.ir_value()], - "redux.sync.max.NaN.f32 $0, $1, 0xffffffff;", - "=f,f", - has_side_effects=False, - is_align_stack=False, - asm_dialect=llvm.AsmDialect.AD_ATT, - ) - return Float32(ret) - - @cute.jit - def smem_fmax(ptr : Pointer, val : Float32): + def smem_fmax(ptr: Pointer, val: Float32): # https://stackoverflow.com/a/72461459 - # Works with canonical NaN which warp_redux_fmax should return + # Works with canonical NaN which warp_redux_sync(kind="fmax") should return llvm.inline_asm( None, [ptr.llvm_ptr, val.ir_value()], @@ -1472,37 +1613,9 @@ class MixedInputFusedMultiHeadAttentionDecode: asm_dialect=llvm.AsmDialect.AD_ATT, ) - @cute.jit - def dsmem_fmax(self, val_ptr: Pointer, val: Float32, mbar_ptr: Pointer): - expect_tx_bytes = Int32(Float32.width // 8) - val_llvm_ptr = self._mapa(val_ptr, 0) - mbar_llvm_ptr = self._mapa(mbar_ptr, 0) - - nvvm.mbarrier_txn( - mbar_llvm_ptr, - expect_tx_bytes.ir_value(), - kind=nvvm.MBarrierTxnKind.ARRIVE_EXPECT_TX, - space=nvvm.MBarrierSpaceKind.CLUSTER, - ) - - llvm.inline_asm( - None, - [val_llvm_ptr, val.ir_value(), mbar_llvm_ptr], - """{\n\t - .reg .pred p;\n\t - setp.lt.s32 p, $1, 0x0; - @p red.async.relaxed.cluster.shared::cluster.mbarrier::complete_tx::bytes.min.u32 [$0], $1, [$2];\n\t - @!p red.async.relaxed.cluster.shared::cluster.mbarrier::complete_tx::bytes.max.s32 [$0], $1, [$2];\n\t - }\n\t""", - "r,r,r", - has_side_effects=True, - is_align_stack=False, - asm_dialect=llvm.AsmDialect.AD_ATT, - ) - @staticmethod @cute.jit - def gmem_fmax(ptr : Pointer, val : Float32): + def gmem_fmax(ptr: Pointer, val: Float32): llvm.inline_asm( None, [ptr.llvm_ptr, val.ir_value()], @@ -1518,53 +1631,26 @@ class MixedInputFusedMultiHeadAttentionDecode: asm_dialect=llvm.AsmDialect.AD_ATT, ) - @cute.jit - def smem_fadd(ptr : Pointer, val : Float32): - # Expensive - llvm.inline_asm( - None, - [ptr.llvm_ptr, val.ir_value()], - "red.relaxed.shared::cta.add.f32 [$0], $1;", - "r,f", - has_side_effects=True, - is_align_stack=False, - asm_dialect=llvm.AsmDialect.AD_ATT, - ) - - @staticmethod - @cute.jit - def gmem_fadd(ptr : Pointer, val : Float32): - llvm.inline_asm( - None, - [ptr.llvm_ptr, val.ir_value()], - "red.relaxed.cluster.global.add.f32 [$0], $1;", - "l,f", - has_side_effects=True, - is_align_stack=False, - asm_dialect=llvm.AsmDialect.AD_ATT, - ) - def run( - batches : int, - seqlen : int, - heads_q : int, - heads_k : int, - headdim : int, - block_scaledim: int, - kv_splits: int, - kv_cluster_dim: int, - q_dtype: Type[cutlass.Numeric], - kv_dtype: Type[cutlass.Numeric], - o_dtype: Type[cutlass.Numeric], - acc_dtype: Type[cutlass.Numeric], - tolerance: float, - scale_q: float, - scale_o: float, - scale_s: float, - warmup_iterations: int, - iterations: int, - skip_ref_check: bool, + batches: int = 1, + seqlen: int = 1024, + heads_q: int = 32, + heads_k: int = 4, + headdim: int = 512, + block_scaledim: int = 512, + kv_splits: int = 0, + q_dtype: Type[cutlass.Numeric] = BFloat16, + kv_dtype: Type[cutlass.Numeric] = Int8, + o_dtype: Type[cutlass.Numeric] = BFloat16, + acc_dtype: Type[cutlass.Numeric] = Float32, + tolerance: float = 0.1, + scale_q: float = 1.0, + scale_o: float = 1.0, + scale_s: float = 0.0, + warmup_iterations: int = 0, + iterations: int = 0, + skip_ref_check: bool = False, use_cold_l2: bool = False, **kwargs, ): @@ -1572,7 +1658,7 @@ def run( print(f"\tbatches: {batches}, seqlen: {seqlen}") print(f"\theads_q: {heads_q}, heads_k: {heads_k}") print(f"\theaddim: {headdim}, block_scaledim: {block_scaledim}") - print(f"\tkv_splits: {kv_splits}, kv_cluster_dim: {kv_cluster_dim}") + print(f"\tkv_splits: {kv_splits}") print(f"\tq_dtype: {q_dtype}") print(f"\tkv_dtype: {kv_dtype}") print(f"\to_dtype: {o_dtype}") @@ -1593,12 +1679,16 @@ def run( # Config Kernel # grouped_heads = heads_q // heads_k + convert_warpgroups = 1 + if headdim == 512 and grouped_heads == 8 and kv_dtype.width == 4: + convert_warpgroups = 4 + elif headdim > 128: + convert_warpgroups = 2 fmha = MixedInputFusedMultiHeadAttentionDecode( headdim=headdim, block_scaledim=block_scaledim, grouped_head_tile=min(cute.round_up(grouped_heads, 8), 32), - dual_convert=(headdim > 128), - deterministic=(kv_cluster_dim == 1), + convert_warpgroups=convert_warpgroups, ) if scale_s == 0.0: # default to 1/sqrt(d) @@ -1610,25 +1700,27 @@ def run( sm_count = hardware_info.get_device_multiprocessor_count() sm_count = 148 if sm_count <= 0 else sm_count grid_yz = batches * heads_k * math.ceil(grouped_heads / fmha.grouped_head_tile) - kv_splits = sm_count // grid_yz # 1 wave + kv_splits = sm_count // grid_yz # 1 wave kv_splits = max(1, kv_splits) if sm_count == 148 and grid_yz == 32: - kv_splits = 9 # 2 waves + kv_splits = 9 # 2 waves print(f"\tauto kv_splits: {kv_splits}") seqlen_q = 1 seqlen_k = seqlen - kv_clusters = kv_splits // kv_cluster_dim problem_shape = (batches, heads_q, heads_k, seqlen_k, headdim) - fmha.can_implement(problem_shape, kv_splits, kv_cluster_dim, q_dtype, kv_dtype, o_dtype, acc_dtype) + fmha.can_implement( + problem_shape, kv_splits, q_dtype, kv_dtype, o_dtype, acc_dtype + ) # # Allocate Tensors # torch.manual_seed(1111) - def create_tensor(shape, dtype, init = True): + + def create_tensor(shape, dtype, init=True): init_type = cutlass.torch.TensorInitType.RANDOM init_config = cutlass.torch.RandomInitConfig(min_val=-2, max_val=2) @@ -1641,10 +1733,14 @@ def run( elif isinstance(init, tuple) or isinstance(init, list): if len(init) == 2: init_type = cutlass.torch.TensorInitType.RANDOM - init_config = cutlass.torch.RandomInitConfig(min_val=init[0], max_val=init[1]) + init_config = cutlass.torch.RandomInitConfig( + min_val=init[0], max_val=init[1] + ) if len(init) == 3: init_type = cutlass.torch.TensorInitType.GAUSSIAN - init_config = cutlass.torch.RandomInitConfig(mean=init[0], std=init[1], scale=init[2]) + init_config = cutlass.torch.RandomInitConfig( + mean=init[0], std=init[1], scale=init[2] + ) f32_torch_tensor = cutlass_torch.create_and_permute_torch_tensor( shape, @@ -1671,20 +1767,26 @@ def run( torch_tensor, ) - qo_shape = (kv_clusters, batches, heads_q, seqlen_q, headdim) - kv_shape = ( batches, heads_k, seqlen_k, headdim) - scale_shape = ( batches, heads_k, seqlen_k, fmha.scaledim) + qo_shape = (kv_splits, batches, heads_q, seqlen_q, headdim) + kv_shape = (batches, heads_k, seqlen_k, headdim) + scale_shape = (batches, heads_k, seqlen_k, fmha.scaledim) - q_ref, q_cute, q_torch = create_tensor(qo_shape[1:], q_dtype, init=[-8,7]) - k_ref, k_cute, k_torch = create_tensor(kv_shape, kv_dtype, init=[-8,7]) - v_ref, v_cute, v_torch = create_tensor(kv_shape, kv_dtype, init=[-8,7]) - k_scale_ref, k_scale_cute, k_scale_torch = create_tensor(scale_shape, q_dtype, init=[-2, 2]) - v_scale_ref, v_scale_cute, v_scale_torch = create_tensor(scale_shape, q_dtype, init=[-2, 2]) + q_ref, q_cute, q_torch = create_tensor(qo_shape[1:], q_dtype, init=[-8, 7]) + k_ref, k_cute, k_torch = create_tensor(kv_shape, kv_dtype, init=[-8, 7]) + v_ref, v_cute, v_torch = create_tensor(kv_shape, kv_dtype, init=[-8, 7]) + k_scale_ref, k_scale_cute, k_scale_torch = create_tensor( + scale_shape, q_dtype, init=[-2, 2] + ) + v_scale_ref, v_scale_cute, v_scale_torch = create_tensor( + scale_shape, q_dtype, init=[-2, 2] + ) _, o_cute, o_torch = create_tensor(qo_shape[1:], o_dtype, init=False) _, m_cute, m_torch = create_tensor(qo_shape[1:-1], acc_dtype, init=-math.inf) _, l_cute, l_torch = create_tensor(qo_shape[1:-1], acc_dtype, init=False) _, o_partial_cute, o_partial_torch = create_tensor(qo_shape, acc_dtype, init=0) - _, m_partial_cute, m_partial_torch = create_tensor(qo_shape[:-1], acc_dtype, init=-math.inf) + _, m_partial_cute, m_partial_torch = create_tensor( + qo_shape[:-1], acc_dtype, init=-math.inf + ) _, l_partial_cute, l_partial_torch = create_tensor(qo_shape[:-1], acc_dtype, init=0) # @@ -1695,7 +1797,6 @@ def run( fmha, problem_shape, kv_splits, - kv_cluster_dim, q_cute.iterator, k_cute.iterator, v_cute.iterator, @@ -1717,13 +1818,17 @@ def run( # # Refcheck # - def run_torch_fmha(q_ref, k_ref, v_ref, k_scale_ref, v_scale_ref, scale_qs=1.0, scale_o=1.0): + def run_torch_fmha( + q_ref, k_ref, v_ref, k_scale_ref, v_scale_ref, scale_qs=1.0, scale_o=1.0 + ): for i in range(0, headdim // block_scaledim): j = i * block_scaledim - k_ref[..., j:j+block_scaledim] *= k_scale_ref[..., i:i+1] - v_ref[..., j:j+block_scaledim] *= v_scale_ref[..., i:i+1] + k_ref[..., j : j + block_scaledim] *= k_scale_ref[..., i : i + 1] + v_ref[..., j : j + block_scaledim] *= v_scale_ref[..., i : i + 1] - with sdpa_kernel([SDPBackend.FLASH_ATTENTION, SDPBackend.MATH], set_priority=True): + with sdpa_kernel( + [SDPBackend.FLASH_ATTENTION, SDPBackend.MATH], set_priority=True + ): o_ref = scaled_dot_product_attention( q_ref, k_ref, @@ -1743,7 +1848,6 @@ def run( compiled_fmha( problem_shape, kv_splits, - kv_cluster_dim, q_cute.iterator, k_cute.iterator, v_cute.iterator, @@ -1760,25 +1864,28 @@ def run( current_stream, ) print("Verifying results...") - o_ref = run_torch_fmha(q_ref, k_ref, v_ref, k_scale_ref, v_scale_ref, scale_qs, scale_o) - torch.testing.assert_close(o_ref, o_torch.float().cpu(), atol=tolerance, rtol=1e-05) + o_ref = run_torch_fmha( + q_ref, k_ref, v_ref, k_scale_ref, v_scale_ref, scale_qs, scale_o + ) + torch.testing.assert_close( + o_ref, o_torch.float().cpu(), atol=tolerance, rtol=1e-05 + ) def generate_tensors(): - _, q_cute, _ = create_tensor(qo_shape[1:], q_dtype, init=[-8,7]) - _, k_cute, _ = create_tensor(kv_shape, kv_dtype, init=[-8,7]) - _, v_cute, _ = create_tensor(kv_shape, kv_dtype, init=[-8,7]) - _, k_scale_cute, _ = create_tensor(scale_shape, q_dtype, init=[-8,7]) - _, v_scale_cute, _ = create_tensor(scale_shape, q_dtype, init=[-8,7]) + _, q_cute, _ = create_tensor(qo_shape[1:], q_dtype, init=[-8, 7]) + _, k_cute, _ = create_tensor(kv_shape, kv_dtype, init=[-8, 7]) + _, v_cute, _ = create_tensor(kv_shape, kv_dtype, init=[-8, 7]) + _, k_scale_cute, _ = create_tensor(scale_shape, q_dtype, init=[-8, 7]) + _, v_scale_cute, _ = create_tensor(scale_shape, q_dtype, init=[-8, 7]) _, o_cute, _ = create_tensor(qo_shape[1:], o_dtype, init=False) _, m_cute, _ = create_tensor(qo_shape[1:-1], acc_dtype, init=-math.inf) _, l_cute, _ = create_tensor(qo_shape[1:-1], acc_dtype, init=False) _, o_partial_cute, _ = create_tensor(qo_shape, acc_dtype, init=0) _, m_partial_cute, _ = create_tensor(qo_shape[:-1], acc_dtype, init=-math.inf) _, l_partial_cute, _ = create_tensor(qo_shape[:-1], acc_dtype, init=0) - return testing.JitArguments( + args = testing.JitArguments( problem_shape, kv_splits, - kv_cluster_dim, q_cute.iterator, k_cute.iterator, v_cute.iterator, @@ -1794,6 +1901,21 @@ def run( scale_o, current_stream, ) + args.add_to_scope( + [ + q_cute, + k_cute, + v_cute, + k_scale_cute, + v_scale_cute, + o_cute, + m_cute, + l_cute, + o_partial_cute, + m_partial_cute, + l_partial_cute, + ] + ) # # Profile @@ -1803,14 +1925,24 @@ def run( q_torch_effective = q_torch.values() if q_torch.is_nested else q_torch k_torch_effective = k_torch.values() if k_torch.is_nested else k_torch v_torch_effective = v_torch.values() if v_torch.is_nested else v_torch - k_scale_torch_effective = k_scale_torch.values() if k_scale_torch.is_nested else k_scale_torch - v_scale_torch_effective = v_scale_torch.values() if v_scale_torch.is_nested else v_scale_torch + k_scale_torch_effective = ( + k_scale_torch.values() if k_scale_torch.is_nested else k_scale_torch + ) + v_scale_torch_effective = ( + v_scale_torch.values() if v_scale_torch.is_nested else v_scale_torch + ) o_torch_effective = o_torch.values() if o_torch.is_nested else o_torch m_torch_effective = m_torch.values() if m_torch.is_nested else m_torch l_torch_effective = l_torch.values() if l_torch.is_nested else l_torch - o_partial_torch_effective = o_partial_torch.values() if o_partial_torch.is_nested else o_partial_torch - m_partial_torch_effective = m_partial_torch.values() if m_partial_torch.is_nested else m_partial_torch - l_partial_torch_effective = l_partial_torch.values() if l_partial_torch.is_nested else l_partial_torch + o_partial_torch_effective = ( + o_partial_torch.values() if o_partial_torch.is_nested else o_partial_torch + ) + m_partial_torch_effective = ( + m_partial_torch.values() if m_partial_torch.is_nested else m_partial_torch + ) + l_partial_torch_effective = ( + l_partial_torch.values() if l_partial_torch.is_nested else l_partial_torch + ) one_workspace_bytes = ( q_torch_effective.numel() * q_torch_effective.element_size() + k_torch_effective.numel() * k_torch_effective.element_size() @@ -1820,9 +1952,12 @@ def run( + o_torch_effective.numel() * o_torch_effective.element_size() + m_torch_effective.numel() * m_torch_effective.element_size() + l_torch_effective.numel() * l_torch_effective.element_size() - + o_partial_torch_effective.numel() * o_partial_torch_effective.element_size() - + m_partial_torch_effective.numel() * m_partial_torch_effective.element_size() - + l_partial_torch_effective.numel() * l_partial_torch_effective.element_size() + + o_partial_torch_effective.numel() + * o_partial_torch_effective.element_size() + + m_partial_torch_effective.numel() + * m_partial_torch_effective.element_size() + + l_partial_torch_effective.numel() + * l_partial_torch_effective.element_size() ) workspace_count = testing.get_workspace_count( one_workspace_bytes, warmup_iterations, iterations @@ -1842,6 +1977,7 @@ def run( return exec_time # Return execution time in microseconds + if __name__ == "__main__": def parse_comma_separated_ints(s: str): @@ -1855,61 +1991,45 @@ if __name__ == "__main__": parser = argparse.ArgumentParser(description="Example of FMHA on Blackwell.") parser.add_argument( - "--batches","--batch","--b", - type=int, - default=1, - help="batch size" + "--batches", "--batch", "--b", type=int, default=1, help="batch size" ) parser.add_argument( - "--seqlen","--seqlen_k","--seq","--s", + "--seqlen", + "--seqlen_k", + "--seq", + "--s", type=int, default=1024, - help="key/value sequence length" + help="key/value sequence length", + ) + + parser.add_argument("--heads_q", "--h_q", type=int, default=32, help="query heads") + + parser.add_argument( + "--heads_k", "--h_k", type=int, default=4, help="key/value heads" ) parser.add_argument( - "--heads_q","--h_q", - type=int, - default=32, - help="query heads" + "--headdim", "--d", type=int, default=512, help="head dimmension" ) parser.add_argument( - "--heads_k","--h_k", - type=int, - default=4, - help="key/value heads" - ) - - parser.add_argument( - "--headdim","--d", + "--block_scaledim", + "--bs", type=int, default=512, - help="head dimmension" + help="headdim per scale factor", ) parser.add_argument( - "--block_scaledim","--bs", - type=int, - default=512, - help="headdim per scale factor" - ) - - parser.add_argument( - "--kv_splits","--splits", + "--kv_splits", + "--splits", type=int, default=0, help="threadblocks per sequence", ) - parser.add_argument( - "--kv_cluster_dim","--cluster", - type=int, - default=1, - help="threadblocks per partial buffer", - ) - parser.add_argument( "--q_dtype", type=cutlass.dtype, @@ -1957,7 +2077,8 @@ if __name__ == "__main__": ) parser.add_argument( - "--scale_s","--scale", + "--scale_s", + "--scale", type=float, default=0.0, help="score (Q*K) scale factor; if zero, defaults to 1/sqrt(D)", @@ -1990,29 +2111,7 @@ if __name__ == "__main__": help="Use circular buffer tensor sets to ensure L2 cold cache", ) - args = parser.parse_args() - - run( - args.batches, - args.seqlen, - args.heads_q, - args.heads_k, - args.headdim, - args.block_scaledim, - args.kv_splits, - args.kv_cluster_dim, - args.q_dtype, - args.kv_dtype, - args.o_dtype, - args.acc_dtype, - args.tolerance, - args.scale_q, - args.scale_o, - args.scale_s, - args.warmup_iterations, - args.iterations, - args.skip_ref_check, - args.use_cold_l2, - ) + kwargs = vars(parser.parse_args()) + run(**kwargs) print("PASS") diff --git a/examples/python/CuTeDSL/blackwell/mixed_input_fmha/mixed_input_fmha_prefill_d256.py b/examples/python/CuTeDSL/blackwell/mixed_input_fmha/mixed_input_fmha_prefill_d256.py new file mode 100644 index 00000000..aaaff2ac --- /dev/null +++ b/examples/python/CuTeDSL/blackwell/mixed_input_fmha/mixed_input_fmha_prefill_d256.py @@ -0,0 +1,2030 @@ +# 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. + +import argparse +import math +import os +import sys +from typing import Type, Tuple, Optional + +import cuda.bindings.driver as cuda +import torch + +import cutlass +import cutlass.cute as cute +import cutlass.cute.nvgpu.tcgen05 as tcgen05 +import cutlass.utils as utils +import cutlass.pipeline as pipeline +from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait +import cutlass.utils.blackwell_helpers as sm100_utils +from cutlass.cute.runtime import from_dlpack +from cutlass.cute.typing import Int32, Int64, Float32 + +if __name__ == "__main__": + current_dir = os.path.dirname(os.path.abspath(__file__)) + sys.path.insert(0, os.path.join(current_dir, "../..")) + +from helpers import fmha_helpers as fmha_utils +from blackwell.mixed_input_fmha import prefill_helpers as prefill_utils + + +class MixedInputFusedMultiHeadAttentionPrefillD256: + def __init__( + self, + scale_granularity: int, + qk_acc_dtype: Type[cutlass.Numeric], + pv_acc_dtype: Type[cutlass.Numeric], + is_persistent: bool, + mask_type: fmha_utils.MaskEnum, + ): + self.qk_acc_dtype = qk_acc_dtype + self.pv_acc_dtype = pv_acc_dtype + self.cta_tiler = (128, 128, 256) + self.qk_mma_tiler = ( + self.cta_tiler[0] * 2, # default 2cta + self.cta_tiler[1], + min(self.cta_tiler[2], 128), # avoid too large GemmK + ) + self.pv_mma_tiler = self.qk_mma_tiler # keep BMM1 & BMM2 at the same pace + self.pv_block_tiler = ( + self.pv_mma_tiler[0] // 2, # default 2cta + self.pv_mma_tiler[1], + self.pv_mma_tiler[2], + ) + self.scale_granularity = scale_granularity + self.iterations_qk = self.cta_tiler[2] // self.qk_mma_tiler[2] + self.iterations_pv = self.cta_tiler[2] // self.pv_mma_tiler[1] + self.cluster_shape_mn = (2, 1) # use 2x1 cluster by default + self.tmem_warp_shape_mn = (4, 1) + self.is_persistent = is_persistent + self.mask_type = mask_type + self.transform_warp_ids = (0, 1, 2, 3, 4, 5, 6, 7) # i8 -> bf16 for kv + self.softmax_warp_ids = (8, 9, 10, 11) # softmax + self.correction_warp_ids = (12, 13, 14, 15) # correction + self.mma_warp_id = 16 # mma + self.load_warp_id = 17 # load + self.empty_warp_ids = (18, 19) # empty + self.num_tmem_alloc_cols = cute.arch.get_max_tmem_alloc_cols("sm_100") + self.tmem_alloc_sync_bar_id = 1 + self.tmem_s_offset = 256 + self.tmem_p_offset = self.tmem_s_offset + self.tmem_o_offset = 0 + self.num_regs_softmax = 256 + self.num_regs_correction = 112 + self.num_regs_other = 32 + self.num_regs_transform = 40 + self.buffer_align_bytes = 1024 + self.threads_per_warp = 32 + self.threads_per_cta = self.threads_per_warp * len( + ( + *self.transform_warp_ids, + *self.softmax_warp_ids, + *self.correction_warp_ids, + self.load_warp_id, + self.mma_warp_id, + *self.empty_warp_ids, + ) + ) + + def _setup_attributes(self): + """Set up configurations and parameters for the FMHA kernel operation. + + This method initializes and configures various attributes required for the + execution of the fused multi-head attention kernel, mainly about the pipeline stages: + + - Sets up staging parameters for Q, K, V inputs and accumulator data + - Configures pipeline stages for softmax, correction, and epilogue operations + """ + + self.q_stage = self.iterations_qk + self.kv_stage = 4 + self.scale_k_stage = self.kv_stage + self.scale_v_stage = self.kv_stage + self.qk_acc_stage = 2 + self.pv_acc_stage = 1 + self.kv_trans_stage = 2 + + @cute.jit + def __call__( + self, + q_iter: cute.Pointer, + k_iter: cute.Pointer, + v_iter: cute.Pointer, + o_iter: cute.Pointer, + scale_k_iter: cute.Pointer, + scale_v_iter: cute.Pointer, + problem_shape: Tuple[Int32, Int32, Int32, Int32, Int32, Int32], + scale_softmax_log2: Float32, + scale_output: Float32, + window_size_left: Optional[Int32], + window_size_right: Optional[Int32], + stream: cuda.CUstream, + ): + self._setup_attributes() + b, s_q, s_k, h_q, h_k, d = problem_shape + h_r = h_q // h_k + self.d_r = self.cta_tiler[2] // self.scale_granularity + # (s, d, ((h_r, h_k), b)) + q_layout = cute.make_layout( + (s_q, d, ((h_r, h_k), b)), + stride=(d, 1, ((d * s_q, d * s_q * h_r), h_r * h_k * s_q * d)), + ) + q = cute.make_tensor(q_iter, q_layout) + # (s, d, ((h_r, h_k), b)), 0-stride for h_r to broadcast + k_layout = cute.make_layout( + (s_k, d, ((h_r, h_k), b)), + stride=(d, 1, ((0, d * s_k), h_k * s_k * d)), + ) + k = cute.make_tensor(k_iter, k_layout) + # (d, s, ((h_r, h_k), b)), 0-stride for h_r to broadcast + v_layout = cute.make_layout( + (d, s_k, ((h_r, h_k), b)), + stride=(1, d, ((0, d * s_k), h_k * s_k * d)), + ) + v = cute.make_tensor(v_iter, v_layout) + # (s, d, ((h_r, h_k), b)) + # set divby for better gmem store vectorization + o_layout = cute.make_layout( + (s_q, d, ((h_r, h_k), b)), + stride=( + cute.assume(d, divby=256), + 1, + ( + ( + cute.assume(d * s_q, divby=256), + cute.assume(d * s_q * h_r, divby=256), + ), + cute.assume(h_r * h_k * s_q * d, divby=256), + ), + ), + ) + o = cute.make_tensor(o_iter, o_layout) + # (d_r * s, ((h_r, h_k), b)) + scale_k_layout = cute.make_layout( + (s_k * self.d_r, ((h_r, h_k), b)), + stride=(1, ((0, self.d_r * s_k), s_k * self.d_r * h_k)), + ) + scale_k = cute.make_tensor(scale_k_iter, scale_k_layout) + # (d_r * s, ((h_r, h_k), b)) + scale_v_layout = cute.make_layout( + (self.d_r * s_k, ((h_r, h_k), b)), + stride=(1, ((0, self.d_r * s_k), s_k * self.d_r * h_k)), + ) + scale_v = cute.make_tensor(scale_v_iter, scale_v_layout) + + self.q_dtype = q.element_type + self.k_dtype = k.element_type + self.v_dtype = v.element_type + self.o_dtype = o.element_type + self.p_dtype = self.q_dtype # pv should has the same dtype + self.tilePlikeFP32 = self.qk_mma_tiler[1] // Float32.width * self.p_dtype.width + self.scale_k_dtype = scale_k.element_type + self.scale_v_dtype = scale_v.element_type + + self.tile_sched_params, grid = fmha_utils.compute_grid( + o.shape, + self.cta_tiler, + self.is_persistent, + ) + + self.q_major_mode = utils.LayoutEnum.from_tensor(q).mma_major_mode() + self.k_major_mode = utils.LayoutEnum.from_tensor(k).mma_major_mode() + self.v_major_mode = utils.LayoutEnum.from_tensor(v).mma_major_mode() + self.o_layout = utils.LayoutEnum.from_tensor(o) + cta_group = tcgen05.CtaGroup.TWO + p_major_mode = tcgen05.OperandMajorMode.K + p_source = tcgen05.OperandSource.TMEM + qk_tiled_mma = sm100_utils.make_trivial_tiled_mma( + self.q_dtype, + self.q_major_mode, + self.k_major_mode, + self.qk_acc_dtype, + cta_group, + self.qk_mma_tiler[:2], + ) + pv_tiled_mma = sm100_utils.make_trivial_tiled_mma( + self.q_dtype, + p_major_mode, + self.v_major_mode, + self.pv_acc_dtype, + cta_group, + self.pv_mma_tiler[:2], + p_source, + ) + self.cluster_shape_mnk = (*self.cluster_shape_mn, 1) + self.cluster_layout_vmnk = cute.tiled_divide( + cute.make_layout(self.cluster_shape_mnk), + (qk_tiled_mma.thr_id.shape,), + ) + self.epi_tile = self.pv_block_tiler[:2] + + q_smem_layout_staged = sm100_utils.make_smem_layout_a( + qk_tiled_mma, + self.qk_mma_tiler, + self.q_dtype, + self.q_stage, + ) + k_smem_layout_staged = sm100_utils.make_smem_layout_b( + qk_tiled_mma, + self.qk_mma_tiler, + self.q_dtype, + self.kv_stage, + ) + k_smem_layout_staged = cute.make_composed_layout( + cute.make_swizzle(0, 4, 3), 0, k_smem_layout_staged.outer + ) + k_trans_smem_layout_staged = sm100_utils.make_smem_layout_b( + qk_tiled_mma, + self.qk_mma_tiler, + self.q_dtype, + self.kv_trans_stage, + ) + p_tmem_layout_staged = sm100_utils.make_smem_layout_a( + pv_tiled_mma, + self.pv_mma_tiler, + self.p_dtype, + self.qk_acc_stage, + ) + p_tmem_layout = cute.select(p_tmem_layout_staged, mode=[0, 1, 2]) + v_smem_layout_staged = sm100_utils.make_smem_layout_b( + pv_tiled_mma, + self.pv_mma_tiler, + self.q_dtype, + self.kv_stage, + ) + v_smem_layout_staged = cute.make_composed_layout( + cute.make_swizzle(0, 4, 3), 0, v_smem_layout_staged.outer + ) + v_trans_smem_layout_staged = sm100_utils.make_smem_layout_b( + pv_tiled_mma, + self.pv_mma_tiler, + self.q_dtype, + self.kv_trans_stage, + ) + scale_k_smem_layout, self.scale_k_tiler, scale_k_s2r_view_layout = ( + prefill_utils.get_scale_smem_layout( + self.scale_granularity, + self.d_r, + self.qk_mma_tiler, + self.k_major_mode, + ) + ) + scale_k_smem_layout_staged = cute.append( + scale_k_smem_layout, + cute.make_layout( + (self.scale_k_stage), + stride=(cute.cosize(scale_k_smem_layout.outer)), + ), + ) + scale_k_s2r_view_layout_staged = cute.append( + scale_k_s2r_view_layout, + cute.make_layout( + (self.scale_k_stage), + stride=(cute.cosize(scale_k_s2r_view_layout)), + ), + ) + scale_v_smem_layout, self.scale_v_tiler, scale_v_s2r_view_layout = ( + prefill_utils.get_scale_smem_layout( + self.scale_granularity, + self.d_r, + self.pv_mma_tiler, + self.v_major_mode, + ) + ) + scale_v_smem_layout_staged = cute.append( + scale_v_smem_layout, + cute.make_layout( + (self.scale_v_stage), + stride=(cute.cosize(scale_v_smem_layout.outer)), + ), + ) + scale_v_s2r_view_layout_staged = cute.append( + scale_v_s2r_view_layout, + cute.make_layout( + (self.scale_v_stage), + stride=(cute.cosize(scale_v_s2r_view_layout)), + ), + ) + + tma_load_q_op = cute.nvgpu.cpasync.CopyBulkTensorTileG2SOp(cta_group) + # For TMA Async, use one cta to sync with corresponding cta only + tma_load_kv_op = cute.nvgpu.cpasync.CopyBulkTensorTileG2SOp( + tcgen05.CtaGroup.ONE + ) + q_smem_layout = cute.select(q_smem_layout_staged, mode=[0, 1, 2]) + tma_atom_q, tma_tensor_q = cute.nvgpu.make_tiled_tma_atom_A( + tma_load_q_op, + q, + q_smem_layout, + self.qk_mma_tiler, + qk_tiled_mma, + self.cluster_layout_vmnk.shape, + ) + # TMA load for K + k_smem_layout = cute.select(k_smem_layout_staged, mode=[0, 1, 2]) + tma_atom_k, tma_tensor_k = cute.nvgpu.make_tiled_tma_atom_B( + tma_load_kv_op, + k, + k_smem_layout, + self.qk_mma_tiler, + qk_tiled_mma, + self.cluster_layout_vmnk.shape, + ) + tma_atom_scale_k, tma_tensor_scale_k = cute.nvgpu.cpasync.make_tiled_tma_atom( + tma_load_kv_op, + scale_k, + scale_k_smem_layout, + (self.scale_k_tiler[0] // 2,), + ) + + # TMA load for V + v_smem_layout = cute.select(v_smem_layout_staged, mode=[0, 1, 2]) + tma_atom_v, tma_tensor_v = cute.nvgpu.make_tiled_tma_atom_B( + tma_load_kv_op, + v, + v_smem_layout, + self.pv_mma_tiler, + pv_tiled_mma, + self.cluster_layout_vmnk.shape, + ) + tma_atom_scale_v, tma_tensor_scale_v = cute.nvgpu.cpasync.make_tiled_tma_atom( + tma_load_kv_op, + scale_v, + scale_v_smem_layout, + self.scale_v_tiler, + ) + + self.tma_copy_q_bytes = cute.size_in_bytes( + self.q_dtype, q_smem_layout + ) * cute.size(qk_tiled_mma.thr_id.shape) + self.tma_copy_kv_bytes = cute.size_in_bytes(self.k_dtype, k_smem_layout) + self.tma_copy_scale_k_bytes = cute.size_in_bytes( + self.scale_k_dtype, scale_k_smem_layout + ) + self.tma_copy_scale_v_bytes = cute.size_in_bytes( + self.scale_v_dtype, scale_v_smem_layout + ) + + @cute.struct + class SharedStorage: + # Pipeline barriers + load_q_mbar_ptr: cute.struct.MemRange[Int64, self.q_stage * 2] + load_kv_mbar_ptr: cute.struct.MemRange[Int64, self.kv_stage * 2] + load_scale_k_mbar_ptr: cute.struct.MemRange[Int64, self.scale_k_stage * 2] + load_scale_v_mbar_ptr: cute.struct.MemRange[Int64, self.scale_v_stage * 2] + dequant_kv_mbar_ptr: cute.struct.MemRange[Int64, self.kv_trans_stage * 2] + mma_s_mbar_ptr: cute.struct.MemRange[Int64, self.qk_acc_stage * 2] + p_mma_mbar_ptr: cute.struct.MemRange[Int64, self.qk_acc_stage * 2] + s_corr_mbar_ptr: cute.struct.MemRange[Int64, self.qk_acc_stage * 2] + sum_mbar_ptr: cute.struct.MemRange[Int64, 2] + mma_o_mbar_ptr: cute.struct.MemRange[Int64, self.pv_acc_stage * 2] + tmem_dealloc_mbar_ptr: Int64 + tmem_holding_buf: Int32 + + self.shared_storage = SharedStorage + + grid = cute.round_up(grid, self.cluster_shape_mnk) + + # Launch the kernel synchronously + self.kernel( + qk_tiled_mma, + pv_tiled_mma, + tma_atom_q, + tma_tensor_q, + tma_atom_k, + tma_tensor_k, + tma_atom_scale_k, + tma_tensor_scale_k, + tma_atom_v, + tma_tensor_v, + tma_atom_scale_v, + tma_tensor_scale_v, + o, + scale_softmax_log2, + scale_output, + window_size_left, + window_size_right, + self.cluster_layout_vmnk, + q_smem_layout_staged, + k_smem_layout_staged, + k_trans_smem_layout_staged, + scale_k_smem_layout_staged, + scale_k_s2r_view_layout_staged, + p_tmem_layout, + v_smem_layout_staged, + v_trans_smem_layout_staged, + scale_v_smem_layout_staged, + scale_v_s2r_view_layout_staged, + self.epi_tile, + self.tile_sched_params, + ).launch( + grid=grid, + block=[self.threads_per_cta, 1, 1], + cluster=self.cluster_shape_mnk, + stream=stream, + min_blocks_per_mp=1, + ) + + @cute.kernel + def kernel( + self, + qk_tiled_mma: cute.TiledMma, + pv_tiled_mma: cute.TiledMma, + tma_atom_q: cute.CopyAtom, + mQ_qdl: cute.Tensor, + tma_atom_k: cute.CopyAtom, + mK_kdl: cute.Tensor, + tma_atom_scale_k: cute.CopyAtom, + mScaleK_kdl: cute.Tensor, + tma_atom_v: cute.CopyAtom, + mV_dkl: cute.Tensor, + tma_atom_scale_v: cute.CopyAtom, + mScaleV_dkl: cute.Tensor, + mO_qdl: cute.Tensor, + scale_softmax_log2: Float32, + scale_output: Float32, + window_size_left: Optional[Int32], + window_size_right: Optional[Int32], + cluster_layout_vmnk: cute.Layout, + q_smem_layout_staged: cute.ComposedLayout, + k_smem_layout_staged: cute.ComposedLayout, + k_trans_smem_layout_staged: cute.ComposedLayout, + scale_k_smem_layout_staged: cute.ComposedLayout, + scale_k_s2r_view_layout_staged: cute.Layout, + p_tmem_layout: cute.ComposedLayout, + v_smem_layout_staged: cute.ComposedLayout, + v_trans_smem_layout_staged: cute.ComposedLayout, + scale_v_smem_layout_staged: cute.ComposedLayout, + scale_v_s2r_view_layout_staged: cute.Layout, + epi_tile: cute.Tile, + tile_sched_params: fmha_utils.FmhaStaticTileSchedulerParams, + ): + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + # Prefetch tma desc + if warp_idx == self.load_warp_id: + cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_q) + cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_k) + cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_v) + cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_scale_k) + cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_scale_v) + bidx, _, _ = cute.arch.block_idx() + mma_tile_coord_v = bidx % cute.size(qk_tiled_mma.thr_id.shape) + cta_rank_in_cluster = cute.arch.make_warp_uniform( + cute.arch.block_idx_in_cluster() + ) + block_in_cluster_coord_vmnk = cluster_layout_vmnk.get_flat_coord( + cta_rank_in_cluster + ) + # Alloc + smem = utils.SmemAllocator() + storage = smem.allocate(self.shared_storage) + + load_q_producer, load_q_consumer = pipeline.PipelineTmaUmma.create( + num_stages=self.q_stage, + producer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, len([self.load_warp_id]) + ), + consumer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, len([self.mma_warp_id]) + ), + tx_count=self.tma_copy_q_bytes, + barrier_storage=storage.load_q_mbar_ptr.data_ptr(), + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ).make_participants() + load_kv_producer, load_kv_consumer = pipeline.PipelineTmaAsync.create( + num_stages=self.kv_stage, + producer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, len([self.load_warp_id]) + ), + consumer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, + len(self.transform_warp_ids) * self.threads_per_warp, + ), + tx_count=self.tma_copy_kv_bytes, + barrier_storage=storage.load_kv_mbar_ptr.data_ptr(), + tidx=0, + defer_sync=True, + ).make_participants() + load_scale_k_producer, load_scale_k_consumer = pipeline.PipelineTmaAsync.create( + num_stages=self.scale_k_stage, + producer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, len([self.load_warp_id]) + ), + consumer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, + len(self.transform_warp_ids) * self.threads_per_warp, + ), + tx_count=self.tma_copy_scale_k_bytes, + barrier_storage=storage.load_scale_k_mbar_ptr.data_ptr(), + tidx=0, + defer_sync=True, + ).make_participants() + load_scale_v_producer, load_scale_v_consumer = pipeline.PipelineTmaAsync.create( + num_stages=self.scale_v_stage, + producer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, len([self.load_warp_id]) + ), + consumer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, + len(self.transform_warp_ids) * self.threads_per_warp, + ), + tx_count=self.tma_copy_scale_v_bytes, + barrier_storage=storage.load_scale_v_mbar_ptr.data_ptr(), + tidx=0, + defer_sync=True, + ).make_participants() + dequant_kv_producer, dequant_kv_consumer = pipeline.PipelineAsyncUmma.create( + num_stages=self.kv_trans_stage, + producer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, + len(self.transform_warp_ids) + * self.threads_per_warp + * self.cluster_shape_mnk[0], + ), + consumer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, len([self.mma_warp_id]) + ), + barrier_storage=storage.dequant_kv_mbar_ptr.data_ptr(), + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ).make_participants() + mma_s_producer, mma_s_consumer = pipeline.PipelineUmmaAsync.create( + num_stages=self.qk_acc_stage, + producer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, len([self.mma_warp_id]) + ), + consumer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, + len(self.softmax_warp_ids) + * self.threads_per_warp + * self.cluster_shape_mnk[0], + ), + barrier_storage=storage.mma_s_mbar_ptr.data_ptr(), + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ).make_participants() + p_mma_producer, p_mma_consumer = pipeline.PipelineAsyncUmma.create( + num_stages=self.qk_acc_stage, + producer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, + len(self.softmax_warp_ids) + * self.threads_per_warp + * self.cluster_shape_mnk[0], + ), + consumer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, len([self.mma_warp_id]) + ), + barrier_storage=storage.p_mma_mbar_ptr.data_ptr(), + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ).make_participants() + s_corr_producer, s_corr_consumer = pipeline.PipelineAsync.create( + num_stages=self.qk_acc_stage, + producer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, + self.threads_per_warp * len(self.softmax_warp_ids), + ), + consumer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, + self.threads_per_warp * len(self.correction_warp_ids), + ), + barrier_storage=storage.s_corr_mbar_ptr.data_ptr(), + defer_sync=True, + ).make_participants() + sum_producer, sum_consumer = pipeline.PipelineAsync.create( + num_stages=1, + producer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, + self.threads_per_warp * len(self.softmax_warp_ids), + ), + consumer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, + self.threads_per_warp * len(self.correction_warp_ids), + ), + barrier_storage=storage.sum_mbar_ptr.data_ptr(), + defer_sync=True, + ).make_participants() + mma_o_producer, mma_o_consumer = pipeline.PipelineUmmaAsync.create( + num_stages=self.pv_acc_stage, + producer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, len([self.mma_warp_id]) + ), + consumer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, + len(self.correction_warp_ids) + * self.threads_per_warp + * self.cluster_shape_mnk[0], + ), + barrier_storage=storage.mma_o_mbar_ptr.data_ptr(), + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ).make_participants() + tmem_alloc_barrier = pipeline.NamedBarrier( + barrier_id=self.tmem_alloc_sync_bar_id, + num_threads=self.threads_per_warp + * len( + (self.mma_warp_id, *self.softmax_warp_ids, *self.correction_warp_ids) + ), + ) + # Tensor memory dealloc barrier init + tmem = utils.TmemAllocator( + storage.tmem_holding_buf, + barrier_for_retrieve=tmem_alloc_barrier, + allocator_warp_id=self.correction_warp_ids[0], + is_two_cta=True, + two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar_ptr, + ) + # Cluster arrive after barrier init + pipeline_init_arrive(cluster_shape_mn=cluster_layout_vmnk, is_relaxed=True) + + sK_trans = smem.allocate_tensor( + element_type=self.q_dtype, + layout=k_trans_smem_layout_staged.outer, + swizzle=k_trans_smem_layout_staged.inner, + byte_alignment=128, + ) + sV_trans_ptr = cute.recast_ptr( + sK_trans.iterator, v_trans_smem_layout_staged.inner + ) + sV_trans = cute.make_tensor(sV_trans_ptr, v_trans_smem_layout_staged.outer) + sQ = smem.allocate_tensor( + element_type=self.q_dtype, + layout=q_smem_layout_staged.outer, + swizzle=q_smem_layout_staged.inner, + byte_alignment=128, + ) + sScaleK = smem.allocate_tensor( + element_type=self.scale_k_dtype, + layout=scale_k_smem_layout_staged.outer, + swizzle=scale_k_smem_layout_staged.inner, + byte_alignment=128, + ) + sScaleK_s2r_view = cute.make_tensor( + sScaleK.iterator, scale_k_s2r_view_layout_staged + ) + sScaleV = smem.allocate_tensor( + element_type=self.scale_v_dtype, + layout=scale_v_smem_layout_staged.outer, + swizzle=scale_v_smem_layout_staged.inner, + byte_alignment=128, + ) + sScaleV_s2r_view = cute.make_tensor( + sScaleV.iterator, scale_v_s2r_view_layout_staged + ) + sK = smem.allocate_tensor( + element_type=self.k_dtype, + layout=k_smem_layout_staged.outer, + swizzle=k_smem_layout_staged.inner, + byte_alignment=128, + ) + sV_ptr = cute.recast_ptr(sK.iterator, v_smem_layout_staged.inner) + sV = cute.make_tensor(sV_ptr, v_smem_layout_staged.outer) + + sSum = smem.allocate_tensor( + element_type=self.qk_acc_dtype, + layout=cute.make_layout(len(self.softmax_warp_ids) * self.threads_per_warp), + byte_alignment=128, + ) + + qk_thr_mma = qk_tiled_mma.get_slice(mma_tile_coord_v) + pv_thr_mma = pv_tiled_mma.get_slice(mma_tile_coord_v) + tSrQ = qk_thr_mma.make_fragment_A(sQ) + tSrK_trans = qk_thr_mma.make_fragment_B(sK_trans) + tOrV_trans = pv_thr_mma.make_fragment_B(sV_trans) + qk_acc_shape = pv_thr_mma.partition_shape_C( + (self.qk_mma_tiler[0], self.qk_mma_tiler[1]) + ) + # (atomV, restM, restN, accStage) + tStS = qk_tiled_mma.make_fragment_C( + cute.append(qk_acc_shape, self.qk_acc_stage) + ) + pv_acc_shape = pv_thr_mma.partition_shape_C( + cute.select(self.pv_mma_tiler, mode=[0, 1]) + ) + # (atomV, restM, restN) + tOtO = pv_thr_mma.make_fragment_C(pv_acc_shape) + tOtO_layout = cute.append( + tOtO.layout, + cute.make_layout( + self.iterations_pv, + stride=self.pv_mma_tiler[1] // self.tmem_warp_shape_mn[1], + ), + ) + tStS = cute.make_tensor(tStS.iterator + self.tmem_s_offset, tStS.layout) + tOtO_staged = cute.make_tensor(tOtO.iterator + self.tmem_o_offset, tOtO_layout) + # Local_tile partition global tensors + q_cta_layout = cute.make_layout( + cute.slice_(cluster_layout_vmnk, (0, 0, None, 0)).shape + ) + # (bM, bK, restM, restK, loopM, loopK, loopL) + gQ_qdl = cute.flat_divide(mQ_qdl, cute.select(self.qk_mma_tiler, mode=[0, 2])) + tSgQ_qdl = qk_thr_mma.partition_A(gQ_qdl) + tQsQ, tQgQ_qdl = cute.nvgpu.cpasync.tma_partition( + tma_atom_q, + block_in_cluster_coord_vmnk[2], + q_cta_layout, + cute.group_modes(sQ, 0, 3), + cute.group_modes(tSgQ_qdl, 0, 3), + ) + kv_cta_layout = cute.make_layout( + cute.slice_(cluster_layout_vmnk, (0, None, 0, 0)).shape + ) + # (bN, bK, loopN, loopK, loopL) + gK_kdl = cute.flat_divide(mK_kdl, cute.select(self.qk_mma_tiler, mode=[1, 2])) + tSgK_kdl = qk_thr_mma.partition_B(gK_kdl) + tKsK, tKgK_kdl = cute.nvgpu.cpasync.tma_partition( + tma_atom_k, + block_in_cluster_coord_vmnk[1], + kv_cta_layout, + cute.group_modes(sK, 0, 3), + cute.group_modes(tSgK_kdl, 0, 3), + ) + # (blk, loopBlk, loopL) + gScaleK_kdl = cute.flat_divide(mScaleK_kdl, self.scale_k_tiler) + # Deal with 2cta + gScaleK_kdl_ = cute.logical_divide(gScaleK_kdl, (self.scale_k_tiler[0] // 2,))[ + (None, mma_tile_coord_v), None, None + ] + tKsScaleK, tKgScaleK_kdl = cute.nvgpu.cpasync.tma_partition( + tma_atom_scale_k, + block_in_cluster_coord_vmnk[1], + kv_cta_layout, + sScaleK, + gScaleK_kdl_, + ) + # (bN, bK, loopN, loopK, loopL) + gV_dkl = cute.flat_divide(mV_dkl, cute.select(self.pv_mma_tiler, mode=[1, 2])) + tOgV_dkl = pv_thr_mma.partition_B(gV_dkl) + tVsV, tVgV_dkl = cute.nvgpu.cpasync.tma_partition( + tma_atom_v, + block_in_cluster_coord_vmnk[1], + kv_cta_layout, + cute.group_modes(sV, 0, 3), + cute.group_modes(tOgV_dkl, 0, 3), + ) + # (bBlk, loopBlk, loopL) + gScaleV_dkl = cute.flat_divide(mScaleV_dkl, self.scale_v_tiler) + tVsScaleV, tVgScaleV_dkl = cute.nvgpu.cpasync.tma_partition( + tma_atom_scale_v, + block_in_cluster_coord_vmnk[1], + kv_cta_layout, + sScaleV, + gScaleV_dkl, + ) + # (bM, bN, loopM, loopN, loopL) + gO_qdl = cute.flat_divide(mO_qdl, cute.select(self.pv_block_tiler, mode=[0, 1])) + cO_qdl = cute.flat_divide( + cute.make_identity_tensor(mO_qdl.shape), + cute.select(self.pv_block_tiler, mode=[0, 1]), + ) + seqlen_q = mQ_qdl.shape[0] + seqlen_k = mK_kdl.shape[0] + tile_sched = fmha_utils.create_fmha_static_tile_scheduler( + tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim() + ) + work_tile = tile_sched.initial_work_tile_info() + # Cluster wait + pipeline_init_wait(cluster_shape_mn=cluster_layout_vmnk) + + # /////////////////////////////////////////////////////////////////////////////// + # Load + # /////////////////////////////////////////////////////////////////////////////// + if warp_idx == self.load_warp_id: + cute.arch.setmaxregister_decrease(self.num_regs_other) + while work_tile.is_valid_tile: + curr_block_coord = work_tile.tile_idx + mma_block_coord = ( + curr_block_coord[0] // cute.size(qk_tiled_mma.thr_id.shape), + curr_block_coord[1], + curr_block_coord[2], + ) + seqlen_kv_loop_steps = fmha_utils.FusedMask.get_trip_count( + self.mask_type, + mma_block_coord, + self.qk_mma_tiler, + seqlen_q, + seqlen_k, + window_size_left, + window_size_right, + ) + # ((atom_v, rest_v), RestK) + tQgQ = tQgQ_qdl[None, mma_block_coord[0], None, mma_block_coord[2]] + # ((atom_v, rest_v), RestN, RestK) + tKgK = tKgK_kdl[None, None, None, mma_block_coord[2]] + tKgScaleK = tKgScaleK_kdl[None, None, mma_block_coord[2]] + # ((atom_v, rest_v), RestN, RestK) + tVgV = tVgV_dkl[None, None, None, mma_block_coord[2]] + tVgScaleV = tVgScaleV_dkl[None, None, mma_block_coord[2]] + load_kv_producer, load_scale_k_producer, load_q_producer = ( + prefill_utils.load_qk( # Q & K0 & ScaleK0 + self.iterations_qk, + kv_step=0, + k_args=(tKgK, tKsK, tma_atom_k, load_kv_producer), + scale_k_args=( + tKgScaleK, + tKsScaleK, + tma_atom_scale_k, + load_scale_k_producer, + ), + q_args=(tQgQ, tQsQ, tma_atom_q, load_q_producer), + ) + ) + for step in cutlass.range(1, seqlen_kv_loop_steps, 1, unroll=1): + load_kv_producer, load_scale_k_producer = ( + prefill_utils.load_qk( # Ki & ScaleKi + self.iterations_qk, + kv_step=step, + k_args=(tKgK, tKsK, tma_atom_k, load_kv_producer), + scale_k_args=( + tKgScaleK, + tKsScaleK, + tma_atom_scale_k, + load_scale_k_producer, + ), + ) + ) + load_kv_producer, load_scale_v_producer = ( + prefill_utils.load_v( # Vi-1 & ScaleVi-1 + self.iterations_pv, + kv_step=step - 1, + v_args=(tVgV, tVsV, tma_atom_v, load_kv_producer), + scale_v_args=( + tVgScaleV, + tVsScaleV, + tma_atom_scale_v, + load_scale_v_producer, + ), + ) + ) + load_kv_producer, load_scale_v_producer = ( + prefill_utils.load_v( # Vend & ScaleVend + self.iterations_pv, + kv_step=seqlen_kv_loop_steps - 1, + v_args=(tVgV, tVsV, tma_atom_v, load_kv_producer), + scale_v_args=( + tVgScaleV, + tVsScaleV, + tma_atom_scale_v, + load_scale_v_producer, + ), + ) + ) + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + load_kv_producer.tail() + load_scale_k_producer.tail() + load_scale_v_producer.tail() + load_q_producer.tail() + + # /////////////////////////////////////////////////////////////////////////////// + # MMA + # /////////////////////////////////////////////////////////////////////////////// + if warp_idx == self.mma_warp_id: + cute.arch.setmaxregister_decrease(self.num_regs_other) + tmem.wait_for_alloc() + while work_tile.is_valid_tile: + curr_block_coord = work_tile.tile_idx + mma_block_coord = ( + curr_block_coord[0] // cute.size(qk_tiled_mma.thr_id.shape), + curr_block_coord[1], + curr_block_coord[2], + ) + seqlen_kv_loop_steps = fmha_utils.FusedMask.get_trip_count( + self.mask_type, + mma_block_coord, + self.qk_mma_tiler, + seqlen_q, + seqlen_k, + window_size_left, + window_size_right, + ) + load_q_releaser = load_q_consumer.clone() + pv_tiled_mma.set(tcgen05.Field.ACCUMULATE, False) + if seqlen_kv_loop_steps > 1: + mma_s_producer, load_q_consumer, dequant_kv_consumer = ( + prefill_utils.mma_qk( # QK0 + self.iterations_qk, + qk_tiled_mma, + (tStS, tSrQ, tSrK_trans), + ( + mma_s_producer, + load_q_consumer, + None, + dequant_kv_consumer, + ), + ) + ) + for i in cutlass.range(1, seqlen_kv_loop_steps - 1, 1, unroll=1): + mma_s_producer, _, dequant_kv_consumer = ( + prefill_utils.mma_qk( # QKi + self.iterations_qk, + qk_tiled_mma, + (tStS, tSrQ, tSrK_trans), + (mma_s_producer, None, None, dequant_kv_consumer), + ) + ) + ( + pv_tiled_mma, + p_mma_consumer, + mma_o_producer, + dequant_kv_consumer, + ) = self.mma_pv( # PVi + (pv_tiled_mma, pv_thr_mma), + (tOtO_staged, tStS, tOrV_trans, p_tmem_layout), + (p_mma_consumer, mma_o_producer, dequant_kv_consumer), + ) + mma_s_producer, _, dequant_kv_consumer = ( + prefill_utils.mma_qk( # QKend needs to release Q + self.iterations_qk, + qk_tiled_mma, + (tStS, tSrQ, tSrK_trans), + ( + mma_s_producer, + None, + load_q_releaser, + dequant_kv_consumer, + ), + ) + ) + ( + pv_tiled_mma, + p_mma_consumer, + mma_o_producer, + dequant_kv_consumer, + ) = self.mma_pv( # PVend-1 + (pv_tiled_mma, pv_thr_mma), + (tOtO_staged, tStS, tOrV_trans, p_tmem_layout), + (p_mma_consumer, mma_o_producer, dequant_kv_consumer), + ) + else: + mma_s_producer, load_q_consumer, dequant_kv_consumer = ( + prefill_utils.mma_qk( # QK0 + self.iterations_qk, + qk_tiled_mma, + (tStS, tSrQ, tSrK_trans), + ( + mma_s_producer, + load_q_consumer, + load_q_releaser, + dequant_kv_consumer, + ), + ) + ) + pv_tiled_mma, p_mma_consumer, mma_o_producer, dequant_kv_consumer = ( + self.mma_pv( # PVend + (pv_tiled_mma, pv_thr_mma), + (tOtO_staged, tStS, tOrV_trans, p_tmem_layout), + (p_mma_consumer, mma_o_producer, dequant_kv_consumer), + ) + ) + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + mma_s_producer.tail() + mma_o_producer.tail() + + # /////////////////////////////////////////////////////////////////////////////// + # Softmax + # /////////////////////////////////////////////////////////////////////////////// + if ( + warp_idx < self.correction_warp_ids[0] + and warp_idx >= self.softmax_warp_ids[0] + ): + cute.arch.setmaxregister_increase(self.num_regs_softmax) + tmem.wait_for_alloc() + while work_tile.is_valid_tile: + curr_block_coord = work_tile.tile_idx + mma_block_coord = ( + curr_block_coord[0] // cute.size(qk_tiled_mma.thr_id.shape), + curr_block_coord[1], + curr_block_coord[2], + ) + seqlen_kv_loop_steps = fmha_utils.FusedMask.get_trip_count( + self.mask_type, + mma_block_coord, + self.qk_mma_tiler, + seqlen_q, + seqlen_k, + window_size_left, + window_size_right, + ) + unmask_steps = fmha_utils.FusedMask.get_unmasked_trip_count( + self.mask_type, + mma_block_coord, + self.qk_mma_tiler, + seqlen_q, + seqlen_k, + window_size_left, + window_size_right, + ) + cS_base = cute.make_identity_tensor( + (self.qk_mma_tiler[0], self.qk_mma_tiler[1]) + ) + cS = cute.domain_offset( + (mma_block_coord[0] * self.qk_mma_tiler[0], 0), cS_base + ) + tScS = qk_thr_mma.partition_C(cS) + row_max = -Float32.inf + row_max_prev = -Float32.inf + row_sum = 0.0 + for step in cutlass.range(seqlen_kv_loop_steps, unroll=1): + cS_iter = cute.domain_offset((0, step * self.qk_mma_tiler[1]), cS) + tScS_iter = qk_thr_mma.partition_C(cS_iter) + # Si -> Pi + ( + row_max, + row_sum, + mma_s_consumer, + p_mma_producer, + s_corr_producer, + ) = self.softmax_step( + (step >= unmask_steps, window_size_left, window_size_right), + ( + row_max_prev, + row_sum, + seqlen_q, + seqlen_k, + scale_softmax_log2, + ), + (tStS, tScS_iter), + (mma_s_consumer, p_mma_producer, s_corr_producer), + ) + row_max_prev = row_max + sum_producer = self.store_sum(row_sum, sSum, sum_producer) + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + p_mma_producer.tail() + s_corr_producer.tail() + + # /////////////////////////////////////////////////////////////////////////////// + # Correction + # /////////////////////////////////////////////////////////////////////////////// + if warp_idx < self.mma_warp_id and warp_idx >= self.correction_warp_ids[0]: + cute.arch.setmaxregister_increase(self.num_regs_correction) + tmem.allocate(self.num_tmem_alloc_cols) + tmem.wait_for_alloc() + tmem_ptr = tmem.retrieve_ptr(self.qk_acc_dtype) + while work_tile.is_valid_tile: + curr_block_coord = work_tile.tile_idx + mma_block_coord = ( + curr_block_coord[0] // cute.size(qk_tiled_mma.thr_id.shape), + curr_block_coord[1], + curr_block_coord[2], + ) + seqlen_kv_loop_steps = fmha_utils.FusedMask.get_trip_count( + self.mask_type, + mma_block_coord, + self.qk_mma_tiler, + seqlen_q, + seqlen_k, + window_size_left, + window_size_right, + ) + gO_staged = gO_qdl[ + None, None, curr_block_coord[0], None, curr_block_coord[2] + ] + cO_staged = cO_qdl[ + None, None, curr_block_coord[0], None, curr_block_coord[2] + ] + cS = cute.make_identity_tensor( + (self.qk_mma_tiler[0], self.qk_mma_tiler[1]) + ) + tScS = qk_thr_mma.partition_C(cS) + # Empty step as the first step is no need for correction + stats_handle = s_corr_consumer.wait_and_advance() + stats_handle.release() + for step in cutlass.range(1, seqlen_kv_loop_steps, 1, unroll=1): + # Oi-1 -> Oi + mma_o_consumer, s_corr_consumer = self.correction_rescale( + scale_softmax_log2, + (s_corr_consumer, tStS, tScS), + (mma_o_consumer, tOtO_staged, cO_staged), + epi_tile, + ) + # O_partial -> O_final + mma_o_consumer, sum_consumer = self.correction_epilog( + (seqlen_q, scale_output), + (sum_consumer, sSum), + (mma_o_consumer, gO_staged, cO_staged, tOtO_staged), + epi_tile, + ) + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + tmem.relinquish_alloc_permit() + tmem.free(tmem_ptr) + + # /////////////////////////////////////////////////////////////////////////////// + # Trans + # /////////////////////////////////////////////////////////////////////////////// + if warp_idx < self.softmax_warp_ids[0]: + cute.arch.setmaxregister_decrease(self.num_regs_transform) + qk_thr_mma_leader_cta = qk_tiled_mma.get_slice(0) + pv_thr_mma_leader_cta = pv_tiled_mma.get_slice(0) + sScaleK_ = qk_thr_mma_leader_cta.partition_B(sScaleK_s2r_view) + sScaleV_ = pv_thr_mma_leader_cta.partition_B(sScaleV_s2r_view) + while work_tile.is_valid_tile: + curr_block_coord = work_tile.tile_idx + mma_block_coord = ( + curr_block_coord[0] // cute.size(qk_tiled_mma.thr_id.shape), + curr_block_coord[1], + curr_block_coord[2], + ) + seqlen_kv_loop_steps = fmha_utils.FusedMask.get_trip_count( + self.mask_type, + mma_block_coord, + self.qk_mma_tiler, + seqlen_q, + seqlen_k, + window_size_left, + window_size_right, + ) + load_kv_consumer, load_scale_k_consumer, dequant_kv_producer = ( + prefill_utils.dequant_k( # K0 + self.iterations_qk, + self.transform_warp_ids, + (self.k_dtype, self.q_dtype), + (sK, sScaleK_, sK_trans), + (load_kv_consumer, load_scale_k_consumer, dequant_kv_producer), + ) + ) + for step in cutlass.range(1, seqlen_kv_loop_steps, 1, unroll=1): + load_kv_consumer, load_scale_k_consumer, dequant_kv_producer = ( + prefill_utils.dequant_k( # Ki + self.iterations_qk, + self.transform_warp_ids, + (self.k_dtype, self.q_dtype), + (sK, sScaleK_, sK_trans), + ( + load_kv_consumer, + load_scale_k_consumer, + dequant_kv_producer, + ), + ) + ) + load_kv_consumer, load_scale_v_consumer, dequant_kv_producer = ( + prefill_utils.dequant_v( # Vi-1 + self.iterations_pv, + self.transform_warp_ids, + (self.v_dtype, self.q_dtype), + (sV, sScaleV_, sV_trans), + ( + load_kv_consumer, + load_scale_v_consumer, + dequant_kv_producer, + ), + ) + ) + load_kv_consumer, load_scale_v_consumer, dequant_kv_producer = ( + prefill_utils.dequant_v( # Vend + self.iterations_pv, + self.transform_warp_ids, + (self.v_dtype, self.q_dtype), + (sV, sScaleV_, sV_trans), + (load_kv_consumer, load_scale_v_consumer, dequant_kv_producer), + ) + ) + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + dequant_kv_producer.tail() + + # /////////////////////////////////////////////////////////////////////////////// + # Empty + # /////////////////////////////////////////////////////////////////////////////// + if warp_idx > self.load_warp_id: + cute.arch.setmaxregister_decrease(self.num_regs_other) + + return + + @cute.jit + def mma_pv( + self, + mma_args: Tuple, + tensor_args: Tuple, + pipeline_args: Tuple, + ): + pv_tiled_mma, pv_thr_mma = mma_args + tOtO_staged, tStS, tOrV_trans, p_tmem_layout = tensor_args + p_mma_consumer, mma_o_producer, dequant_kv_consumer = pipeline_args + cta_rank_in_cluster = cute.arch.make_warp_uniform( + cute.arch.block_idx_in_cluster() + ) + is_leader_cta = cta_rank_in_cluster % 2 == 0 + if is_leader_cta: + p_handle = p_mma_consumer.wait_and_advance() + o_handle = mma_o_producer.acquire_and_advance() + pv_whether_acc = pv_tiled_mma.get(tcgen05.Field.ACCUMULATE) + for iter in cutlass.range(self.iterations_pv, unroll=1): + v_trans_handle = dequant_kv_consumer.wait_and_advance() + pv_tiled_mma.set(tcgen05.Field.ACCUMULATE, pv_whether_acc) + tOtO_slice = tOtO_staged[None, None, None, iter] + tStS_slice = tStS[None, None, None, p_handle.index] + tP = cute.make_tensor(tStS_slice.iterator, p_tmem_layout.outer) + tOrP = pv_thr_mma.make_fragment_A(tP) + tOrP_slice = cute.make_tensor( + cute.recast_ptr(tStS_slice.iterator, dtype=self.p_dtype), + tOrP.layout, + ) + tOrV_trans_slice = tOrV_trans[None, None, None, v_trans_handle.index] + num_kphases = cute.size(tOrV_trans_slice, mode=[2]) + for kphase_idx in cutlass.range(num_kphases, unroll_full=True): + kphase_coord = (None, None, kphase_idx) + cute.gemm( + pv_tiled_mma, + tOtO_slice, + tOrP_slice[kphase_coord], + tOrV_trans_slice[kphase_coord], + tOtO_slice, + ) + pv_tiled_mma.set(tcgen05.Field.ACCUMULATE, True) + v_trans_handle.release() + o_handle.commit() + p_handle.release() + return pv_tiled_mma, p_mma_consumer, mma_o_producer, dequant_kv_consumer + + @cute.jit + def softmax_step( + self, + mask_args: Tuple, + value_args: Tuple, + tensor_args: Tuple, + pipeline_args: Tuple, + ) -> Tuple[Float32, Float32, pipeline.PipelineConsumer, pipeline.PipelineProducer]: + need_apply_mask, window_size_left, window_size_right = mask_args + row_max, row_sum, seqlen_q, seqlen_k, scale_softmax_log2 = value_args + tStS, tScS = tensor_args + mma_s_consumer, p_mma_producer, s_corr_producer = pipeline_args + tidx, _, _ = cute.arch.thread_idx() + thread_idx = tidx % (self.threads_per_warp * len(self.softmax_warp_ids)) + s_handle = mma_s_consumer.wait_and_advance() + tStS_slice = tStS[(None, None), 0, 0, s_handle.index] + tScS_slice = tScS[(None, None), 0, 0] + tmem_load_atom = cute.make_copy_atom( + tcgen05.Ld32x32bOp(tcgen05.Repetition(32)), self.qk_acc_dtype + ) + tmem_tiled_load = tcgen05.make_tmem_copy(tmem_load_atom, tStS_slice) + thr_load = tmem_tiled_load.get_slice(thread_idx) + tTMEM_LOADtS = thr_load.partition_S(tStS_slice) + tTMEM_LOADcS = thr_load.partition_D(tScS_slice) + tTMEM_LOADrS = cute.make_rmem_tensor(tTMEM_LOADcS.shape, self.qk_acc_dtype) + cute.copy(tmem_tiled_load, tTMEM_LOADtS, tTMEM_LOADrS) + cute.arch.fence_view_async_tmem_load() + s_handle.release() + if need_apply_mask: + fmha_utils.FusedMask.apply_mask( + self.mask_type, + tTMEM_LOADrS, + tTMEM_LOADcS, + seqlen_q, + seqlen_k, + window_size_left, + window_size_right, + ) + old_row_max = row_max + row_max = tTMEM_LOADrS.load().reduce(cute.ReductionOp.MAX, row_max, 0) + row_max_safe = row_max + if row_max == -cutlass.Float32.inf: + row_max_safe = 0.0 + + stats_handle = s_corr_producer.acquire_and_advance() + stats_layout = cute.composition( + tStS_slice.layout, cute.make_layout((tStS_slice.shape[0], 2)) + ) + stats_c_layout = cute.composition( + tScS_slice.layout, cute.make_layout((tScS_slice.shape[0], 2)) + ) + tOtStats = cute.make_tensor( + tStS_slice.iterator + self.tilePlikeFP32, stats_layout + ) + tOcStats = cute.make_tensor(tScS_slice.iterator, stats_c_layout) + tmem_store_stats_atom = cute.make_copy_atom( + tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(2)), + self.qk_acc_dtype, + ) + tiled_tmem_store_stats = tcgen05.make_tmem_copy(tmem_store_stats_atom, tOtStats) + thr_tmem_store_stats = tiled_tmem_store_stats.get_slice(thread_idx) + tTMEM_STOREcStats = thr_tmem_store_stats.partition_S(tOcStats) + tTMEM_STORErStats = cute.make_rmem_tensor( + tTMEM_STOREcStats.shape, self.qk_acc_dtype + ) + tTMEM_STORErStats[0] = old_row_max + tTMEM_STORErStats[1] = row_max_safe + tTMEM_STOREtStats = thr_tmem_store_stats.partition_D(tOtStats) + cute.copy(tiled_tmem_store_stats, tTMEM_STORErStats, tTMEM_STOREtStats) + cute.arch.fence_view_async_tmem_store() + stats_handle.commit() + + scale = scale_softmax_log2 + minus_row_max_scale = (0.0 - row_max_safe) * scale + tTMEM_STORErP = cute.make_rmem_tensor(tTMEM_LOADrS.shape, self.p_dtype) + for k in cutlass.range(cute.size(tTMEM_LOADrS), vectorize=True): + tTMEM_LOADrS[k] = tTMEM_LOADrS[k] * scale + minus_row_max_scale + tTMEM_LOADrS[k] = cute.math.exp2(tTMEM_LOADrS[k], fastmath=True) + s_vec = tTMEM_LOADrS.load() + tTMEM_STORErP.store(s_vec.to(self.p_dtype)) + + p_handle = p_mma_producer.acquire_and_advance() + tmem_store_atom = cute.make_copy_atom( + tcgen05.St32x32bOp(tcgen05.Repetition(32)), self.qk_acc_dtype + ) + tilePlikeFP32 = tStS_slice.shape[1] // Float32.width * self.p_dtype.width + tStS_P_layout = cute.composition( + tStS_slice.layout, cute.make_layout((tStS_slice.shape[0], tilePlikeFP32)) + ) + tStS_P = cute.make_tensor(tStS_slice.iterator, tStS_P_layout) + tScS_P_layout = cute.composition( + tScS_slice.layout, cute.make_layout((tScS_slice.shape[0], tilePlikeFP32)) + ) + tScS_P = cute.make_tensor(tScS_slice.iterator, tScS_P_layout) + tmem_tiled_store = tcgen05.make_tmem_copy(tmem_store_atom, tStS_P) + thr_store = tmem_tiled_store.get_slice(thread_idx) + tTMEM_STOREtP = thr_store.partition_D(tStS_P) + tTMEM_STOREcS = thr_store.partition_S(tScS_P) + tTMEM_STORErP_ = cute.make_tensor( + cute.recast_ptr(tTMEM_STORErP.iterator, dtype=self.qk_acc_dtype), + tTMEM_STOREcS.shape, + ) + cute.copy(tmem_tiled_store, tTMEM_STORErP_, tTMEM_STOREtP) + cute.arch.fence_view_async_tmem_store() + + p_handle.commit() + acc_scale_ = scale * (old_row_max - row_max_safe) + acc_scale = cute.math.exp2(acc_scale_, fastmath=True) * 0.5 + # TODO: calc row sum with TensorSSA + row_sum *= acc_scale + local_row_sum_0 = (row_sum, row_sum) + local_row_sum_1 = (0.0, 0.0) + local_row_sum_2 = (0.0, 0.0) + local_row_sum_3 = (0.0, 0.0) + reduction_unroll = 4 + frg_tile = cute.size(tTMEM_LOADrS) // reduction_unroll + tTMEM_LOADrS_frg = cute.logical_divide(tTMEM_LOADrS, cute.make_layout(frg_tile)) + for j in cutlass.range_constexpr(0, cute.size(tTMEM_LOADrS_frg, mode=[0]), 2): + local_row_sum_0 = cute.arch.add_packed_f32x2( + local_row_sum_0, (tTMEM_LOADrS_frg[j, 0], tTMEM_LOADrS_frg[j + 1, 0]) + ) + local_row_sum_1 = cute.arch.add_packed_f32x2( + local_row_sum_1, (tTMEM_LOADrS_frg[j, 1], tTMEM_LOADrS_frg[j + 1, 1]) + ) + local_row_sum_2 = cute.arch.add_packed_f32x2( + local_row_sum_2, (tTMEM_LOADrS_frg[j, 2], tTMEM_LOADrS_frg[j + 1, 2]) + ) + local_row_sum_3 = cute.arch.add_packed_f32x2( + local_row_sum_3, (tTMEM_LOADrS_frg[j, 3], tTMEM_LOADrS_frg[j + 1, 3]) + ) + local_row_sum_0 = cute.arch.add_packed_f32x2(local_row_sum_0, local_row_sum_1) + local_row_sum_2 = cute.arch.add_packed_f32x2(local_row_sum_2, local_row_sum_3) + local_row_sum_0 = cute.arch.add_packed_f32x2(local_row_sum_0, local_row_sum_2) + row_sum = local_row_sum_0[0] + local_row_sum_0[1] + return row_max, row_sum, mma_s_consumer, p_mma_producer, s_corr_producer + + @cute.jit + def correction_rescale( + self, + scale_softmax_log2: Float32, + stats_args: tuple, + o_args: tuple, + epi_tile: cute.Tile, + ) -> pipeline.PipelineConsumer: + (s_corr_consumer, tStS, tScS) = stats_args + (mma_o_consumer, tOtO_staged, cO_staged) = o_args + tidx, _, _ = cute.arch.thread_idx() + thread_idx = tidx % (self.threads_per_warp * len(self.softmax_warp_ids)) + + stats_handle = s_corr_consumer.wait_and_advance() + tStS_slice = tStS[(None, None), 0, 0, stats_handle.index] + tScS_slice = tScS[(None, None), 0, 0] + stats_layout = cute.composition( + tStS_slice.layout, cute.make_layout((tStS_slice.shape[0], 2)) + ) + stats_c_layout = cute.composition( + tScS_slice.layout, cute.make_layout((tScS_slice.shape[0], 2)) + ) + tOtStats = cute.make_tensor( + tStS_slice.iterator + self.tilePlikeFP32, stats_layout + ) + tOcStats = cute.make_tensor(tScS_slice.iterator, stats_c_layout) + tmem_load_stats_atom = cute.make_copy_atom( + tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(2)), + self.qk_acc_dtype, + ) + tiled_tmem_load_stats = tcgen05.make_tmem_copy(tmem_load_stats_atom, tOtStats) + thr_tmem_load_stats = tiled_tmem_load_stats.get_slice(thread_idx) + tTMEM_LOADtStats = thr_tmem_load_stats.partition_S(tOtStats) + tTMEM_LOADcStats = thr_tmem_load_stats.partition_D(tOcStats) + tTMEM_LOADrStats = cute.make_rmem_tensor( + tTMEM_LOADcStats.shape, self.qk_acc_dtype + ) + cute.copy(tiled_tmem_load_stats, tTMEM_LOADtStats, tTMEM_LOADrStats) + + scale = scale_softmax_log2 * (tTMEM_LOADrStats[0] - tTMEM_LOADrStats[1]) + scale = cute.math.exp2(scale, fastmath=True) + stats_handle.release() + o_handle = mma_o_consumer.wait_and_advance() + for iter in cutlass.range(self.iterations_pv, unroll_full=True): + tOtO = tOtO_staged[(None, None), 0, 0, iter] + cO = cO_staged[None, None, iter] + tOtO_epi = cute.zipped_divide(tOtO, epi_tile) + cO_epi = cute.zipped_divide(cO, epi_tile) + tmem_load_atom = cute.make_copy_atom( + tcgen05.Ld32x32bOp(tcgen05.Repetition(16)), + self.pv_acc_dtype, + ) + tmem_tiled_load = tcgen05.make_tmem_copy(tmem_load_atom, tOtO_epi) + thr_load = tmem_tiled_load.get_slice(thread_idx) + tmem_store_atom = cute.make_copy_atom( + tcgen05.St32x32bOp(tcgen05.Repetition(16)), + self.pv_acc_dtype, + ) + tmem_store_atom = tcgen05.make_tmem_copy(tmem_store_atom, tOtO_epi) + thr_store = tmem_store_atom.get_slice(thread_idx) + tTMEM_LOADtO = thr_load.partition_S(tOtO_epi) + tTMEM_LOADcO = thr_load.partition_D(cO_epi) + tTMEM_STOREtO = thr_store.partition_D(tOtO_epi) + tTMrO = cute.make_rmem_tensor_like( + cute.append( + cute.make_layout(tTMEM_LOADcO[None, 0, 0].shape), + cute.make_layout( + 2, stride=cute.size(tTMEM_LOADcO[None, 0, 0].shape) + ), + ), + self.pv_acc_dtype, + ) + tTMEM_LOADtO_0 = tTMEM_LOADtO[None, 0, 0] + cute.copy(tmem_tiled_load, tTMEM_LOADtO_0, tTMrO[None, 0]) + iter_num = cute.size(tTMEM_LOADtO, mode=[1]) + for i in cutlass.range(1, iter_num, unroll_full=True): + tTMEM_LOADtO_i = tTMEM_LOADtO[None, i, 0] + cute.copy(tmem_tiled_load, tTMEM_LOADtO_i, tTMrO[None, i % 2]) + for j in cutlass.range( + cute.size(tTMrO, mode=[0]), unroll_full=True, vectorize=True + ): + tTMrO[j, (i - 1) % 2] = tTMrO[j, (i - 1) % 2] * scale + tTMEM_STOREtO_prev_i = tTMEM_STOREtO[None, i - 1, 0] + cute.copy( + tmem_store_atom, tTMrO[None, (i - 1) % 2], tTMEM_STOREtO_prev_i + ) + + for j in cutlass.range( + cute.size(tTMrO, mode=[0]), unroll_full=True, vectorize=True + ): + tTMrO[j, (iter_num - 1) % 2] = tTMrO[j, (iter_num - 1) % 2] * scale + cute.copy( + tmem_store_atom, + tTMrO[None, (iter_num - 1) % 2], + tTMEM_STOREtO[None, iter_num - 1, 0], + ) + cute.arch.fence_view_async_tmem_store() + o_handle.release() + return mma_o_consumer, s_corr_consumer + + @cute.jit + def correction_epilog( + self, + value_args: Tuple, + sum_args: Tuple, + o_args: Tuple, + epi_tile: cute.Tile, + ) -> Tuple[pipeline.PipelineConsumer, pipeline.PipelineProducer]: + (seqlen_q, scale_output) = value_args + (sum_consumer, sSum) = sum_args + (mma_o_consumer, gO_staged, cO_staged, tOtO_staged) = o_args + tidx, _, _ = cute.arch.thread_idx() + thread_idx = tidx % (self.threads_per_warp * len(self.softmax_warp_ids)) + sum_handle = sum_consumer.wait_and_advance() + row_sum = sSum[thread_idx] + cute.arch.fence_view_async_shared() + sum_handle.release() + scale = scale_output / row_sum + o_handle = mma_o_consumer.wait_and_advance() + for iter in cutlass.range(self.iterations_pv): + gO = gO_staged[None, None, iter] + cO = cO_staged[None, None, iter] + tOtO = tOtO_staged[(None, None), 0, 0, iter] + tOtO_epi = cute.zipped_divide(tOtO, epi_tile) + cO_epi = cute.zipped_divide(cO, epi_tile) + gO_epi = cute.zipped_divide(gO, epi_tile) + tidx, _, _ = cute.arch.thread_idx() + thread_idx = tidx % (self.threads_per_warp * len(self.softmax_warp_ids)) + tmem_copy_atom = cute.make_copy_atom( + tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(32)), self.pv_acc_dtype + ) + tiled_tmem_load = tcgen05.make_tmem_copy(tmem_copy_atom, tOtO_epi) + thr_tmem_load = tiled_tmem_load.get_slice(thread_idx) + tTMEM_LOADtO = thr_tmem_load.partition_S(tOtO_epi) + tTMEM_LOADgO = thr_tmem_load.partition_D(gO_epi) + tTMEM_LOADcO = thr_tmem_load.partition_D(cO_epi) + for i in cutlass.range(cute.size(tTMEM_LOADtO, mode=[1]), unroll_full=True): + tTMEM_LOADtO_i = tTMEM_LOADtO[None, i, 0] + tTMEM_LOADgO_i = tTMEM_LOADgO[None, i, 0] + tTMEM_LOADcO_i = tTMEM_LOADcO[None, i, 0] + tTMrO = cute.make_rmem_tensor( + tTMEM_LOADcO[None, 0, i].shape, self.pv_acc_dtype + ) + cute.copy(tiled_tmem_load, tTMEM_LOADtO_i, tTMrO) + for j in cutlass.range( + cute.size(tTMrO), unroll_full=True, vectorize=True + ): + tTMrO[j] = tTMrO[j] * scale + tSMrO = cute.make_rmem_tensor(tTMrO.shape, self.o_dtype) + o_vec = tTMrO.load() + tSMrO.store(o_vec.to(self.o_dtype)) + if cute.elem_less(tTMEM_LOADcO_i[0][0], seqlen_q): + cute.autovec_copy(tSMrO, tTMEM_LOADgO_i) + o_handle.release() + return mma_o_consumer, sum_consumer + + def store_sum(self, row_sum, sSum, sum_producer): + tidx, _, _ = cute.arch.thread_idx() + thread_idx = tidx % (self.threads_per_warp * len(self.softmax_warp_ids)) + sum_handle = sum_producer.acquire_and_advance() + sSum[thread_idx] = row_sum + cute.arch.fence_view_async_shared() + sum_handle.commit() + return sum_producer + + +def run( + q_shape: Tuple[int, int, int, int], + k_shape: Tuple[int, int, int, int], + q_dtype: Type[cutlass.Numeric], + kv_dtype: Type[cutlass.Numeric], + o_dtype: Type[cutlass.Numeric], + scale_dtype: Type[cutlass.Numeric], + scale_granularity: int, + qk_acc_dtype: Type[cutlass.Numeric], + pv_acc_dtype: Type[cutlass.Numeric], + is_persistent: bool, + is_causal: bool, + scale_q: float, + scale_k: float, + scale_v: float, + inv_scale_o: float, + scale_softmax: float, + tolerance: float, + warmup_iterations: int, + iterations: int, + skip_ref_check: bool, + use_cold_l2: bool = False, + **kwargs, +): + print(f"Running Blackwell SM100 Mixed Input FMHA Prefill D256 test with:") + print(f" q_shape: {q_shape}") + print(f" k_shape: {k_shape}") + print(f" q_dtype: {q_dtype}") + print(f" kv_dtype: {kv_dtype}") + print(f" o_dtype: {o_dtype}") + print(f" scale_dtype: {scale_dtype}") + print(f" scale_granularity: {scale_granularity}") + print(f" qk_acc_dtype: {qk_acc_dtype}") + print(f" pv_acc_dtype: {pv_acc_dtype}") + print(f" is_persistent: {is_persistent}") + print(f" is_causal: {is_causal}") + print(f" scale_q: {scale_q}") + print(f" scale_k: {scale_k}") + print(f" scale_v: {scale_v}") + print(f" inv_scale_o: {inv_scale_o}") + print(f" scale_softmax: {scale_softmax}") + print(f" tolerance: {tolerance}") + print(f" warmup_iterations: {warmup_iterations}") + print(f" iterations: {iterations}") + print(f" skip_ref_check: {skip_ref_check}") + print(f" use_cold_l2: {use_cold_l2}") + import cutlass.torch as cutlass_torch + + # Unpack parameters + b, h_q, s_q, d = q_shape + b_, h_k, s_k, d_ = k_shape + window_size_left, window_size_right = None, None + if is_causal: + window_size_right = 0 + + if b != b_: + raise ValueError("q & k must have the same batch size") + + if d != d_: + raise ValueError("q & k must have the same head dimension") + + if d not in {256}: + raise ValueError("head dimension must be 256") + + if d % scale_granularity != 0: + raise ValueError("head dimension must be divisible by scale_granularity") + + if scale_granularity not in {128, 256}: + raise ValueError("scale_granularity must be 128, 256") + + if h_q % h_k != 0: + raise ValueError("h_q must be divisible by h_k") + + if isinstance(s_q, tuple) and len(s_q) != b: + raise ValueError("variable_seqlen s_q must have the length of batch size") + if isinstance(s_k, tuple) and len(s_k) != b: + raise ValueError("variable_seqlen s_k must have the length of batch size") + + if q_dtype not in {cutlass.BFloat16}: + raise ValueError("in_dtype must be BFloat16") + + if o_dtype not in {cutlass.BFloat16}: + raise ValueError("o_dtype must be BFloat16") + + if kv_dtype not in {cutlass.Int8}: + raise ValueError("kv_dtype must be Int8") + + if qk_acc_dtype not in {cutlass.Float32}: + raise ValueError("qk_acc_dtype must be Float32") + + if pv_acc_dtype not in {cutlass.Float32}: + raise ValueError("pv_acc_dtype must be Float32") + + h_r = h_q // h_k + + if not torch.cuda.is_available(): + raise RuntimeError("GPU is required to run this example!") + + torch.manual_seed(1111) + + def create_tensor(shape, dtype): + f32_torch_tensor = cutlass_torch.create_and_permute_torch_tensor( + shape, + torch.float32, + permute_order=None, + init_type=cutlass.torch.TensorInitType.RANDOM, + init_config=cutlass.torch.RandomInitConfig( + min_val=-2 if dtype.is_float or dtype.signed else 0, max_val=2 + ), + ) + + _, torch_tensor = cutlass_torch.cute_tensor_like( + f32_torch_tensor, + dtype, + is_dynamic_layout=True, + assumed_align=32, + ) + + # Create dtype cute tensor with offset (gpu) + cute_tensor = from_dlpack(torch_tensor, assumed_align=128) + cute_tensor.element_type = dtype + + return ( + f32_torch_tensor, + cute_tensor, + torch_tensor, + ) + + scale_shape = (b, h_k, s_k, d // scale_granularity) + + q_ref, q_tensor, q_torch = create_tensor(q_shape, q_dtype) + k_ref, k_tensor, k_torch = create_tensor(k_shape, kv_dtype) + v_ref, v_tensor, v_torch = create_tensor(k_shape, kv_dtype) + o_ref, o_tensor, o_torch = create_tensor(q_shape, o_dtype) + scale_k_ref, scale_k_tensor, scale_k_torch = create_tensor(scale_shape, scale_dtype) + scale_v_ref, scale_v_tensor, scale_v_torch = create_tensor(scale_shape, scale_dtype) + + mask_type = fmha_utils.MaskEnum.WINDOW_MASK_INFERENCE + if is_causal: + mask_type = fmha_utils.MaskEnum.WINDOW_MASK_INFERENCE + else: + if s_k % 128 != 0: + mask_type = fmha_utils.MaskEnum.RESIDUAL_MASK + + fmha = MixedInputFusedMultiHeadAttentionPrefillD256( + scale_granularity, + qk_acc_dtype, + pv_acc_dtype, + is_persistent, + mask_type, + ) + + # Initialize Stream + current_stream = cutlass_torch.default_stream() + + if scale_softmax == 0.0: # default to 1/sqrt(d) + scale_softmax = 1.0 / math.sqrt(d) + log2_e = math.log2( + math.exp(1.0) + ) # gpu uses exp2 for perf concerns, we need an extra factor 'log2_e' here + + scale_softmax = scale_q * scale_k * scale_softmax + scale_softmax_log2 = scale_softmax * log2_e + scale_output = scale_v * inv_scale_o + problem_size = (b, s_q, s_k, h_q, h_k, d) + compiled_fmha = cute.compile( + fmha, + q_tensor.iterator, + k_tensor.iterator, + v_tensor.iterator, + o_tensor.iterator, + scale_k_tensor.iterator, + scale_v_tensor.iterator, + problem_size, + scale_softmax_log2, + scale_output, + window_size_left if window_size_left is None else Int32(window_size_left), + window_size_right if window_size_right is None else Int32(window_size_right), + current_stream, + options=f"--opt-level 2", + ) + + def run_torch_fmha( + q, k, v, scale_k, scale_v, scale_softmax=1.0, scale_output=1.0, is_causal=False + ): + h_q = q.shape[1] + h_k = k.shape[1] + if not h_q == h_k: + repeat_factor = h_q // h_k + k = k.repeat_interleave(repeat_factor, dim=1) + v = v.repeat_interleave(repeat_factor, dim=1) + scale_k = scale_k.repeat_interleave(repeat_factor, dim=1) + scale_v = scale_v.repeat_interleave(repeat_factor, dim=1) + scale_k = ( + scale_k.unsqueeze(-1) + .repeat(1, 1, 1, 1, k.shape[3] // scale_k.shape[3]) + .reshape(k.shape) + ) + scale_v = ( + scale_v.unsqueeze(-1) + .repeat(1, 1, 1, 1, v.shape[3] // scale_v.shape[3]) + .reshape(v.shape) + ) + batch = q.shape[0] + ref_list = [] + for batch_idx in range(batch): + q_i = q[batch_idx] + k_i = k[batch_idx] + v_i = v[batch_idx] + scale_k_i = scale_k[batch_idx] + scale_v_i = scale_v[batch_idx] + s_i = torch.einsum("hqd,hkd->hqk", q_i, k_i * scale_k_i) * scale_softmax + s_q = q_i.shape[1] + s_k = k_i.shape[1] + if is_causal: + q_coords = torch.arange(0, s_q).view(-1, 1) + k_coords = torch.arange(0, s_k).view(1, -1) + _mask = k_coords > q_coords + s_k - s_q + s_i = s_i.masked_fill(_mask, -torch.inf) + p_i = s_i.softmax(dim=-1) + ref_i = torch.einsum("hqk,hkd->hqd", p_i, v_i * scale_v_i) * scale_output + ref_list.append(ref_i) + ref = torch.stack(ref_list) + return ref + + if not skip_ref_check: + # Execute kernel once for reference checking + compiled_fmha( + q_tensor.iterator, + k_tensor.iterator, + v_tensor.iterator, + o_tensor.iterator, + scale_k_tensor.iterator, + scale_v_tensor.iterator, + problem_size, + scale_softmax_log2, + scale_output, + window_size_left if window_size_left is None else Int32(window_size_left), + ( + window_size_right + if window_size_right is None + else Int32(window_size_right) + ), + current_stream, + ) + print("Verifying results...") + o_ref = run_torch_fmha( + q_ref, + k_ref, + v_ref, + scale_k_ref, + scale_v_ref, + scale_softmax, + scale_output, + is_causal, + ) + + # convert o back to f32 for comparison + o_fp32, o_fp32_torch = cutlass_torch.cute_tensor_like( + torch.empty(*o_torch.shape, dtype=torch.float32), + Float32, + is_dynamic_layout=True, + assumed_align=16, + ) + cute.testing.convert(o_tensor, o_fp32) + o_result = o_fp32_torch.cpu() + torch.testing.assert_close(o_ref, o_result, atol=tolerance, rtol=1e-05) + + print("Results verified successfully!") + + +if __name__ == "__main__": + + def parse_comma_separated_ints(s: str): + try: + return tuple(int(x.strip()) for x in s.split(",")) + except ValueError: + raise argparse.ArgumentTypeError( + "Invalid format. Expected comma-separated integers." + ) + + parser = argparse.ArgumentParser(description="Example of FMHA on Blackwell.") + + parser.add_argument( + "--q_dtype", + type=cutlass.dtype, + default=cutlass.BFloat16, + help="Input data type", + ) + + parser.add_argument( + "--kv_dtype", + type=cutlass.dtype, + default=cutlass.Int8, + help="Input data type", + ) + + parser.add_argument( + "--o_dtype", + type=cutlass.dtype, + default=cutlass.BFloat16, + help="Output data type", + ) + + parser.add_argument( + "--scale_dtype", + type=cutlass.dtype, + default=cutlass.BFloat16, + help="Scale data type", + ) + + parser.add_argument( + "--scale_granularity", + type=int, + default=256, + help="Scale granularity", + ) + + parser.add_argument( + "--qk_acc_dtype", + type=cutlass.dtype, + default=Float32, + help="QK accumulator data type", + ) + + parser.add_argument( + "--pv_acc_dtype", + type=cutlass.dtype, + default=Float32, + help="PV accumulator data type", + ) + + parser.add_argument( + "--is_persistent", + action="store_true", + help="Is persistent", + ) + + parser.add_argument( + "--is_causal", + action="store_true", + help="Whether to use casual mask", + ) + + parser.add_argument( + "--q_shape", + type=parse_comma_separated_ints, + default=(1, 8, 256, 256), + help="Shape of Q (B, H, S_q, D)", + ) + + parser.add_argument( + "--k_shape", + type=parse_comma_separated_ints, + default=(1, 8, 256, 256), + help="Shape of K (B, H_k, S_k, D)", + ) + + parser.add_argument( + "--scale_q", + type=float, + default=1.0, + help="Scaling factors to dequantize Q", + ) + + parser.add_argument( + "--scale_k", + type=float, + default=1.0, + help="Scaling factors to dequantize K", + ) + + parser.add_argument( + "--scale_v", + type=float, + default=1.0, + help="Scaling factors to dequantize V", + ) + + parser.add_argument( + "--inv_scale_o", + type=float, + default=1.0, + help="Scaling factor to quantize O", + ) + + parser.add_argument( + "--scale_softmax", + type=float, + default=0.0, + help="Scaling factor to scale S (i.e. Q*K); if zero, defaults to 1/sqrt(D)", + ) + + parser.add_argument( + "--tolerance", type=float, default=1e-01, help="Tolerance for validation" + ) + + parser.add_argument( + "--warmup_iterations", + type=int, + default=0, + help="Number of iterations for warmup", + ) + + parser.add_argument( + "--iterations", + type=int, + default=1, + help="Number of iterations after warmup", + ) + + parser.add_argument( + "--skip_ref_check", + action="store_true", + help="Skip reference check", + ) + + parser.add_argument( + "--use_cold_l2", + action="store_true", + default=False, + help="Use circular buffer tensor sets to ensure L2 cold cache", + ) + + args = parser.parse_args() + + if len(args.q_shape) != 4: + parser.error("--q_shape must contain exactly 4 values") + + if len(args.k_shape) != 4: + parser.error("--k_shape must contain exactly 4 values") + + if not torch.cuda.is_available(): + raise RuntimeError("GPU is required to run this example!") + + run( + args.q_shape, + args.k_shape, + args.q_dtype, + args.kv_dtype, + args.o_dtype, + args.scale_dtype, + args.scale_granularity, + args.qk_acc_dtype, + args.pv_acc_dtype, + args.is_persistent, + args.is_causal, + args.scale_q, + args.scale_k, + args.scale_v, + args.inv_scale_o, + args.scale_softmax, + args.tolerance, + args.warmup_iterations, + args.iterations, + args.skip_ref_check, + args.use_cold_l2, + ) + + print("PASS") diff --git a/examples/python/CuTeDSL/blackwell/mixed_input_fmha/mixed_input_fmha_prefill_d512.py b/examples/python/CuTeDSL/blackwell/mixed_input_fmha/mixed_input_fmha_prefill_d512.py new file mode 100644 index 00000000..874ffc49 --- /dev/null +++ b/examples/python/CuTeDSL/blackwell/mixed_input_fmha/mixed_input_fmha_prefill_d512.py @@ -0,0 +1,2172 @@ +# 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. + +import argparse +import math +import os +import sys + +from typing import Type, Tuple, Optional + +import cuda.bindings.driver as cuda +import torch + +import cutlass +import cutlass.cute as cute +import cutlass.cute.nvgpu.tcgen05 as tcgen05 +import cutlass.utils as utils +import cutlass.pipeline as pipeline +from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait +import cutlass.utils.blackwell_helpers as sm100_utils +from cutlass.cute.runtime import from_dlpack +from cutlass.cute.typing import Int32, Int64, Float32 + + +if __name__ == "__main__": + current_dir = os.path.dirname(os.path.abspath(__file__)) + sys.path.insert(0, os.path.join(current_dir, "../..")) + +from helpers import fmha_helpers as fmha_utils +from blackwell.mixed_input_fmha import prefill_helpers as prefill_utils + + +class MixedInputFusedMultiHeadAttentionPrefillD512: + def __init__( + self, + scale_granularity: int, + qk_acc_dtype: Type[cutlass.Numeric], + pv_acc_dtype: Type[cutlass.Numeric], + is_persistent: bool, + mask_type: fmha_utils.MaskEnum, + ): + self.qk_acc_dtype = qk_acc_dtype + self.pv_acc_dtype = pv_acc_dtype + self.cta_tiler = (128, 128, 512) # seq_q, seq_k, d + self.qk_mma_tiler = ( + self.cta_tiler[0] * 2, # default 2cta + self.cta_tiler[1], + min(self.cta_tiler[2], 128), + ) + self.pv_mma_tiler = self.qk_mma_tiler # keep BMM1 & BMM2 at the same pace + self.pv_block_tiler = ( + self.pv_mma_tiler[0] // 2, # default 2cta + self.pv_mma_tiler[1], + self.pv_mma_tiler[2], + ) + self.scale_granularity = scale_granularity + self.iterations_qk = self.cta_tiler[2] // self.qk_mma_tiler[2] + self.iterations_pv = self.cta_tiler[2] // self.pv_mma_tiler[1] + self.cluster_shape_mn = (2, 1) # use 2x1 cluster by default + self.tmem_warp_shape_mn = (4, 1) + self.is_persistent = is_persistent + self.mask_type = mask_type + self.transform_warp_ids = (0, 1, 2, 3, 4, 5, 6, 7) # i8 -> bf16 for kv + self.softmax_warp_ids = (8, 9, 10, 11) # softmax + correction + self.mma_warp_id = 12 # mma + self.load_warp_id = 13 # load + self.empty_warp_ids = (14, 15) # empty + self.num_tmem_alloc_cols = cute.arch.get_max_tmem_alloc_cols("sm_100") + self.tmem_alloc_sync_bar_id = 1 + self.tmem_s_offset = 128 + self.tmem_o_offset = 0 + self.num_regs_softmax = 256 + self.num_regs_other = 32 + self.num_regs_transform = 112 + self.buffer_align_bytes = 1024 + self.threads_per_warp = 32 + self.threads_per_cta = self.threads_per_warp * len( + ( + *self.transform_warp_ids, + *self.softmax_warp_ids, + self.load_warp_id, + self.mma_warp_id, + *self.empty_warp_ids, + ) + ) + + def _setup_attributes(self): + """Set up configurations and parameters for the FMHA kernel operation. + + This method initializes and configures various attributes required for the + execution of the fused multi-head attention kernel, mainly about the pipeline stages: + + - Sets up staging parameters for Q, K, V inputs and accumulator data + - Configures pipeline stages for softmax, correction, and epilogue operations + """ + + self.q_stage = self.iterations_qk + self.kv_stage = 4 + self.scale_k_stage = 1 + self.scale_v_stage = 1 + self.qk_acc_stage = 1 + self.pv_acc_stage = 1 + self.swap_stage = 1 + self.kv_trans_stage = 2 + + @cute.jit + def __call__( + self, + q_iter: cute.Pointer, + k_iter: cute.Pointer, + v_iter: cute.Pointer, + o_iter: cute.Pointer, + scale_k_iter: cute.Pointer, + scale_v_iter: cute.Pointer, + problem_shape: Tuple[Int32, Int32, Int32, Int32, Int32, Int32], + scale_softmax_log2: Float32, + scale_output: Float32, + window_size_left: Optional[Int32], + window_size_right: Optional[Int32], + stream: cuda.CUstream, + ): + self._setup_attributes() + b, s_q, s_k, h_q, h_k, d = problem_shape + h_r = h_q // h_k + self.d_r = self.cta_tiler[2] // self.scale_granularity + # (s, d, ((h_r, h_k), b)) + q_layout = cute.make_layout( + (s_q, d, ((h_r, h_k), b)), + stride=(d, 1, ((d * s_q, d * s_q * h_r), h_r * h_k * s_q * d)), + ) + q = cute.make_tensor(q_iter, q_layout) + # (s, d, ((h_r, h_k), b)), 0-stride for h_r to broadcast + k_layout = cute.make_layout( + (s_k, d, ((h_r, h_k), b)), + stride=(d, 1, ((0, d * s_k), h_k * s_k * d)), + ) + k = cute.make_tensor(k_iter, k_layout) + # (d, s, ((h_r, h_k), b)), 0-stride for h_r to broadcast + v_layout = cute.make_layout( + (d, s_k, ((h_r, h_k), b)), + stride=(1, d, ((0, d * s_k), h_k * s_k * d)), + ) + v = cute.make_tensor(v_iter, v_layout) + # (s, d, ((h_r, h_k), b)) + # set divby for better gmem store vectorization + o_layout = cute.make_layout( + (s_q, d, ((h_r, h_k), b)), + stride=( + cute.assume(d, divby=256), + 1, + ( + ( + cute.assume(d * s_q, divby=256), + cute.assume(d * s_q * h_r, divby=256), + ), + cute.assume(h_r * h_k * s_q * d, divby=256), + ), + ), + ) + o = cute.make_tensor(o_iter, o_layout) + # (d_r * s, ((h_r, h_k), b)) + scale_k_layout = cute.make_layout( + (s_k * self.d_r, ((h_r, h_k), b)), + stride=(1, ((0, self.d_r * s_k), s_k * self.d_r * h_k)), + ) + scale_k = cute.make_tensor(scale_k_iter, scale_k_layout) + # (d_r * s, ((h_r, h_k), b)) + scale_v_layout = cute.make_layout( + (self.d_r * s_k, ((h_r, h_k), b)), + stride=(1, ((0, self.d_r * s_k), s_k * self.d_r * h_k)), + ) + scale_v = cute.make_tensor(scale_v_iter, scale_v_layout) + self.q_dtype = q.element_type + self.k_dtype = k.element_type + self.v_dtype = v.element_type + self.o_dtype = o.element_type + self.p_dtype = self.q_dtype # pv should has the same dtype + self.scale_k_dtype = scale_k.element_type + self.scale_v_dtype = scale_v.element_type + self.tile_sched_params, grid = fmha_utils.compute_grid( + o.shape, + self.cta_tiler, + self.is_persistent, + ) + self.q_major_mode = utils.LayoutEnum.from_tensor(q).mma_major_mode() + self.k_major_mode = utils.LayoutEnum.from_tensor(k).mma_major_mode() + self.v_major_mode = utils.LayoutEnum.from_tensor(v).mma_major_mode() + self.o_layout = utils.LayoutEnum.from_tensor(o) + cta_group = tcgen05.CtaGroup.TWO + p_major_mode = tcgen05.OperandMajorMode.K + qk_tiled_mma = sm100_utils.make_trivial_tiled_mma( + self.q_dtype, + self.q_major_mode, + self.k_major_mode, + self.qk_acc_dtype, + cta_group, + self.qk_mma_tiler[:2], + ) + pv_tiled_mma = sm100_utils.make_trivial_tiled_mma( + self.q_dtype, + p_major_mode, + self.v_major_mode, + self.pv_acc_dtype, + cta_group, + self.pv_mma_tiler[:2], + ) + self.cluster_shape_mnk = (*self.cluster_shape_mn, 1) + self.cluster_layout_vmnk = cute.tiled_divide( + cute.make_layout(self.cluster_shape_mnk), + (qk_tiled_mma.thr_id.shape,), + ) + self.epi_tile = self.pv_block_tiler[:2] + q_smem_layout_staged = sm100_utils.make_smem_layout_a( + qk_tiled_mma, + self.qk_mma_tiler, + self.q_dtype, + self.q_stage, + ) + k_smem_layout_staged = sm100_utils.make_smem_layout_b( + qk_tiled_mma, + self.qk_mma_tiler, + self.q_dtype, + self.kv_stage, + ) + k_smem_layout_staged = cute.make_composed_layout( + cute.make_swizzle(0, 4, 3), 0, k_smem_layout_staged.outer + ) + k_trans_smem_layout_staged = sm100_utils.make_smem_layout_b( + qk_tiled_mma, + self.qk_mma_tiler, + self.q_dtype, + self.kv_trans_stage, + ) + p_smem_layout_staged = sm100_utils.make_smem_layout_a( + pv_tiled_mma, + self.pv_mma_tiler, + self.p_dtype, + self.qk_acc_stage, + ) + v_smem_layout_staged = sm100_utils.make_smem_layout_b( + pv_tiled_mma, + self.pv_mma_tiler, + self.q_dtype, + self.kv_stage, + ) + v_smem_layout_staged = cute.make_composed_layout( + cute.make_swizzle(0, 4, 3), 0, v_smem_layout_staged.outer + ) + v_trans_smem_layout_staged = sm100_utils.make_smem_layout_b( + pv_tiled_mma, + self.pv_mma_tiler, + self.q_dtype, + self.kv_trans_stage, + ) + scale_k_smem_layout, self.scale_k_tiler, scale_k_s2r_view_layout = ( + prefill_utils.get_scale_smem_layout( + self.scale_granularity, + self.d_r, + self.qk_mma_tiler, + self.k_major_mode, + ) + ) + scale_k_smem_layout_staged = cute.append( + scale_k_smem_layout, + cute.make_layout( + (self.scale_k_stage), + stride=(cute.cosize(scale_k_smem_layout.outer)), + ), + ) + scale_k_s2r_view_layout_staged = cute.append( + scale_k_s2r_view_layout, + cute.make_layout( + (self.scale_k_stage), + stride=(cute.cosize(scale_k_s2r_view_layout)), + ), + ) + scale_v_smem_layout, self.scale_v_tiler, scale_v_s2r_view_layout = ( + prefill_utils.get_scale_smem_layout( + self.scale_granularity, + self.d_r, + self.pv_mma_tiler, + self.v_major_mode, + ) + ) + scale_v_smem_layout_staged = cute.append( + scale_v_smem_layout, + cute.make_layout( + (self.scale_v_stage), + stride=(cute.cosize(scale_v_smem_layout.outer)), + ), + ) + scale_v_s2r_view_layout_staged = cute.append( + scale_v_s2r_view_layout, + cute.make_layout( + (self.scale_v_stage), + stride=(cute.cosize(scale_v_s2r_view_layout)), + ), + ) + tma_load_q_op = cute.nvgpu.cpasync.CopyBulkTensorTileG2SOp(cta_group) + # For TMA Async, use one cta to sync with corresponding cta only + tma_load_kv_op = cute.nvgpu.cpasync.CopyBulkTensorTileG2SOp( + tcgen05.CtaGroup.ONE + ) + q_smem_layout = cute.select(q_smem_layout_staged, mode=[0, 1, 2]) + tma_atom_q, tma_tensor_q = cute.nvgpu.make_tiled_tma_atom_A( + tma_load_q_op, + q, + q_smem_layout, + self.qk_mma_tiler, + qk_tiled_mma, + self.cluster_layout_vmnk.shape, + ) + # TMA load for K + k_smem_layout = cute.select(k_smem_layout_staged, mode=[0, 1, 2]) + tma_atom_k, tma_tensor_k = cute.nvgpu.make_tiled_tma_atom_B( + tma_load_kv_op, + k, + k_smem_layout, + self.qk_mma_tiler, + qk_tiled_mma, + self.cluster_layout_vmnk.shape, + ) + tma_atom_scale_k, tma_tensor_scale_k = cute.nvgpu.cpasync.make_tiled_tma_atom( + tma_load_kv_op, + scale_k, + scale_k_smem_layout, + (self.scale_k_tiler[0] // 2,), + ) + # TMA load for V + v_smem_layout = cute.select(v_smem_layout_staged, mode=[0, 1, 2]) + tma_atom_v, tma_tensor_v = cute.nvgpu.make_tiled_tma_atom_B( + tma_load_kv_op, + v, + v_smem_layout, + self.pv_mma_tiler, + pv_tiled_mma, + self.cluster_layout_vmnk.shape, + ) + tma_atom_scale_v, tma_tensor_scale_v = cute.nvgpu.cpasync.make_tiled_tma_atom( + tma_load_kv_op, + scale_v, + scale_v_smem_layout, + self.scale_v_tiler, + ) + self.tma_copy_q_bytes = cute.size_in_bytes( + self.q_dtype, q_smem_layout + ) * cute.size(qk_tiled_mma.thr_id.shape) + self.tma_copy_kv_bytes = cute.size_in_bytes(self.k_dtype, k_smem_layout) + self.tma_copy_scale_k_bytes = cute.size_in_bytes( + self.scale_k_dtype, scale_k_smem_layout + ) + self.tma_copy_scale_v_bytes = cute.size_in_bytes( + self.scale_v_dtype, scale_v_smem_layout + ) + + @cute.struct + class SharedStorage: + # Pipeline barriers + load_q_mbar_ptr: cute.struct.MemRange[Int64, self.q_stage * 2] + load_kv_mbar_ptr: cute.struct.MemRange[Int64, self.kv_stage * 2] + load_scale_k_mbar_ptr: cute.struct.MemRange[Int64, self.scale_k_stage * 2] + load_scale_v_mbar_ptr: cute.struct.MemRange[Int64, self.scale_v_stage * 2] + dequant_kv_mbar_ptr: cute.struct.MemRange[Int64, self.kv_trans_stage * 2] + mma_s_mbar_ptr: cute.struct.MemRange[Int64, self.qk_acc_stage * 2] + p_mma_mbar_ptr: cute.struct.MemRange[Int64, self.qk_acc_stage * 2] + mma_o_mbar_ptr: cute.struct.MemRange[Int64, self.pv_acc_stage * 2] + swap_mbar_ptr: cute.struct.MemRange[Int64, self.swap_stage * 2] + tmem_dealloc_mbar_ptr: Int64 + tmem_holding_buf: Int32 + + self.shared_storage = SharedStorage + + grid = cute.round_up(grid, self.cluster_shape_mnk) + + # Launch the kernel synchronously + self.kernel( + qk_tiled_mma, + pv_tiled_mma, + tma_atom_q, + tma_tensor_q, + tma_atom_k, + tma_tensor_k, + tma_atom_scale_k, + tma_tensor_scale_k, + tma_atom_v, + tma_tensor_v, + tma_atom_scale_v, + tma_tensor_scale_v, + o, + scale_softmax_log2, + scale_output, + window_size_left, + window_size_right, + self.cluster_layout_vmnk, + q_smem_layout_staged, + k_smem_layout_staged, + k_trans_smem_layout_staged, + scale_k_smem_layout_staged, + scale_k_s2r_view_layout_staged, + p_smem_layout_staged, + v_smem_layout_staged, + v_trans_smem_layout_staged, + scale_v_smem_layout_staged, + scale_v_s2r_view_layout_staged, + self.epi_tile, + self.tile_sched_params, + ).launch( + grid=grid, + block=[self.threads_per_cta, 1, 1], + cluster=self.cluster_shape_mnk, + stream=stream, + min_blocks_per_mp=1, + ) + + @cute.kernel + def kernel( + self, + qk_tiled_mma: cute.TiledMma, + pv_tiled_mma: cute.TiledMma, + tma_atom_q: cute.CopyAtom, + mQ_qdl: cute.Tensor, + tma_atom_k: cute.CopyAtom, + mK_kdl: cute.Tensor, + tma_atom_scale_k: cute.CopyAtom, + mScaleK_kdl: cute.Tensor, + tma_atom_v: cute.CopyAtom, + mV_dkl: cute.Tensor, + tma_atom_scale_v: cute.CopyAtom, + mScaleV_dkl: cute.Tensor, + mO_qdl: cute.Tensor, + scale_softmax_log2: Float32, + scale_output: Float32, + window_size_left: Optional[Int32], + window_size_right: Optional[Int32], + cluster_layout_vmnk: cute.Layout, + q_smem_layout_staged: cute.ComposedLayout, + k_smem_layout_staged: cute.ComposedLayout, + k_trans_smem_layout_staged: cute.ComposedLayout, + scale_k_smem_layout_staged: cute.ComposedLayout, + scale_k_s2r_view_layout_staged: cute.Layout, + p_smem_layout_staged: cute.ComposedLayout, + v_smem_layout_staged: cute.ComposedLayout, + v_trans_smem_layout_staged: cute.ComposedLayout, + scale_v_smem_layout_staged: cute.ComposedLayout, + scale_v_s2r_view_layout_staged: cute.Layout, + epi_tile: cute.Tile, + tile_sched_params: fmha_utils.FmhaStaticTileSchedulerParams, + ): + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + # Prefetch tma desc + if warp_idx == self.load_warp_id: + cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_q) + cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_k) + cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_v) + cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_scale_k) + cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_scale_v) + bidx, _, _ = cute.arch.block_idx() + mma_tile_coord_v = bidx % cute.size(qk_tiled_mma.thr_id.shape) + cta_rank_in_cluster = cute.arch.make_warp_uniform( + cute.arch.block_idx_in_cluster() + ) + block_in_cluster_coord_vmnk = cluster_layout_vmnk.get_flat_coord( + cta_rank_in_cluster + ) + # Alloc + smem = utils.SmemAllocator() + storage = smem.allocate(self.shared_storage) + + load_q_producer, load_q_consumer = pipeline.PipelineTmaUmma.create( + num_stages=self.q_stage, + producer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, len([self.load_warp_id]) + ), + consumer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, len([self.mma_warp_id]) + ), + tx_count=self.tma_copy_q_bytes, + barrier_storage=storage.load_q_mbar_ptr.data_ptr(), + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ).make_participants() + load_kv_producer, load_kv_consumer = pipeline.PipelineTmaAsync.create( + num_stages=self.kv_stage, + producer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, len([self.load_warp_id]) + ), + consumer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, + len(self.transform_warp_ids) * self.threads_per_warp, + ), + tx_count=self.tma_copy_kv_bytes, + barrier_storage=storage.load_kv_mbar_ptr.data_ptr(), + tidx=0, + defer_sync=True, + ).make_participants() + load_scale_k_producer, load_scale_k_consumer = pipeline.PipelineTmaAsync.create( + num_stages=self.scale_k_stage, + producer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, len([self.load_warp_id]) + ), + consumer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, + len(self.transform_warp_ids) * self.threads_per_warp, + ), + tx_count=self.tma_copy_scale_k_bytes, + barrier_storage=storage.load_scale_k_mbar_ptr.data_ptr(), + defer_sync=True, + tidx=0, + ).make_participants() + load_scale_v_producer, load_scale_v_consumer = pipeline.PipelineTmaAsync.create( + num_stages=self.scale_v_stage, + producer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, len([self.load_warp_id]) + ), + consumer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, + len(self.transform_warp_ids) * self.threads_per_warp, + ), + tx_count=self.tma_copy_scale_v_bytes, + barrier_storage=storage.load_scale_v_mbar_ptr.data_ptr(), + defer_sync=True, + tidx=0, + ).make_participants() + dequant_kv_producer, dequant_kv_consumer = pipeline.PipelineAsyncUmma.create( + num_stages=self.kv_trans_stage, + producer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, + len(self.transform_warp_ids) + * self.threads_per_warp + * self.cluster_shape_mnk[0], + ), + consumer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, len([self.mma_warp_id]) + ), + barrier_storage=storage.dequant_kv_mbar_ptr.data_ptr(), + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ).make_participants() + mma_s_producer, mma_s_consumer = pipeline.PipelineUmmaAsync.create( + num_stages=self.qk_acc_stage, + producer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, len([self.mma_warp_id]) + ), + consumer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, + len(self.softmax_warp_ids) + * self.threads_per_warp + * self.cluster_shape_mnk[0], + ), + barrier_storage=storage.mma_s_mbar_ptr.data_ptr(), + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ).make_participants() + p_mma_producer, p_mma_consumer = pipeline.PipelineAsyncUmma.create( + num_stages=self.qk_acc_stage, + producer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, + len(self.softmax_warp_ids) + * self.threads_per_warp + * self.cluster_shape_mnk[0], + ), + consumer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, len([self.mma_warp_id]) + ), + barrier_storage=storage.p_mma_mbar_ptr.data_ptr(), + cta_layout_vmnk=cluster_layout_vmnk, + ).make_participants() + mma_o_producer, mma_o_consumer = pipeline.PipelineUmmaAsync.create( + num_stages=self.pv_acc_stage, + producer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, len([self.mma_warp_id]) + ), + consumer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, + len(self.softmax_warp_ids) + * self.threads_per_warp + * self.cluster_shape_mnk[0], + ), + barrier_storage=storage.mma_o_mbar_ptr.data_ptr(), + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ).make_participants() + swap_producer, swap_consumer = pipeline.PipelineUmmaAsync.create( + num_stages=self.swap_stage, + producer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, len([self.mma_warp_id]) + ), + consumer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, + len(self.softmax_warp_ids) + * self.threads_per_warp + * self.cluster_shape_mnk[0], + ), + barrier_storage=storage.swap_mbar_ptr.data_ptr(), + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ).make_participants() + tmem_alloc_barrier = pipeline.NamedBarrier( + barrier_id=self.tmem_alloc_sync_bar_id, + num_threads=self.threads_per_warp + * len((self.mma_warp_id, *self.softmax_warp_ids)), + ) + # Tensor memory dealloc barrier init + tmem = utils.TmemAllocator( + storage.tmem_holding_buf, + barrier_for_retrieve=tmem_alloc_barrier, + allocator_warp_id=self.softmax_warp_ids[0], + is_two_cta=True, + two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar_ptr, + ) + # Cluster arrive after barrier init + pipeline_init_arrive(cluster_shape_mn=cluster_layout_vmnk, is_relaxed=True) + + sK_trans = smem.allocate_tensor( + element_type=self.q_dtype, + layout=k_trans_smem_layout_staged.outer, + swizzle=k_trans_smem_layout_staged.inner, + byte_alignment=128, + ) + sV_trans_ptr = cute.recast_ptr( + sK_trans.iterator, v_trans_smem_layout_staged.inner + ) + sV_trans = cute.make_tensor(sV_trans_ptr, v_trans_smem_layout_staged.outer) + sQ = smem.allocate_tensor( + element_type=self.q_dtype, + layout=q_smem_layout_staged.outer, + swizzle=q_smem_layout_staged.inner, + byte_alignment=128, + ) + sScaleK = smem.allocate_tensor( + element_type=self.scale_k_dtype, + layout=scale_k_smem_layout_staged.outer, + swizzle=scale_k_smem_layout_staged.inner, + byte_alignment=128, + ) + sScaleK_s2r_view = cute.make_tensor( + sScaleK.iterator, scale_k_s2r_view_layout_staged + ) + sScaleV = smem.allocate_tensor( + element_type=self.scale_v_dtype, + layout=scale_v_smem_layout_staged.outer, + swizzle=scale_v_smem_layout_staged.inner, + byte_alignment=128, + ) + sScaleV_s2r_view = cute.make_tensor( + sScaleV.iterator, scale_v_s2r_view_layout_staged + ) + sP = smem.allocate_tensor( + element_type=self.p_dtype, + layout=p_smem_layout_staged.outer, + swizzle=p_smem_layout_staged.inner, + byte_alignment=128, + ) + sK = smem.allocate_tensor( + element_type=self.k_dtype, + layout=k_smem_layout_staged.outer, + swizzle=k_smem_layout_staged.inner, + byte_alignment=128, + ) + sV_ptr = cute.recast_ptr(sK.iterator, v_smem_layout_staged.inner) + sV = cute.make_tensor(sV_ptr, v_smem_layout_staged.outer) + qk_thr_mma = qk_tiled_mma.get_slice(mma_tile_coord_v) + pv_thr_mma = pv_tiled_mma.get_slice(mma_tile_coord_v) + tSrQ = qk_thr_mma.make_fragment_A(sQ) + tOrP = qk_thr_mma.make_fragment_A(sP) + tSrK_trans = qk_thr_mma.make_fragment_B(sK_trans) + tOrV_trans = pv_thr_mma.make_fragment_B(sV_trans) + qk_acc_shape = pv_thr_mma.partition_shape_C( + (self.qk_mma_tiler[0], self.qk_mma_tiler[1]) + ) + # (atomV, restM, restN, accStage) + tStS = qk_tiled_mma.make_fragment_C( + cute.append(qk_acc_shape, self.qk_acc_stage) + ) + pv_acc_shape = pv_thr_mma.partition_shape_C( + cute.select(self.pv_mma_tiler, mode=[0, 1]) + ) + # (atomV, restM, restN) + tOtO = pv_thr_mma.make_fragment_C(pv_acc_shape) + tOtO_layout = cute.append( + tOtO.layout, + cute.make_layout( + self.iterations_pv, + stride=self.pv_mma_tiler[1] // self.tmem_warp_shape_mn[1], + ), + ) + tStS = cute.make_tensor(tStS.iterator + self.tmem_s_offset, tStS.layout) + tOtO_staged = cute.make_tensor(tOtO.iterator + self.tmem_o_offset, tOtO_layout) + # Local_tile partition global tensors + q_cta_layout = cute.make_layout( + cute.slice_(cluster_layout_vmnk, (0, 0, None, 0)).shape + ) + # (bM, bK, restM, restK, loopM, loopK, loopL) + gQ_qdl = cute.flat_divide(mQ_qdl, cute.select(self.qk_mma_tiler, mode=[0, 2])) + tSgQ_qdl = qk_thr_mma.partition_A(gQ_qdl) + tQsQ, tQgQ_qdl = cute.nvgpu.cpasync.tma_partition( + tma_atom_q, + block_in_cluster_coord_vmnk[2], + q_cta_layout, + cute.group_modes(sQ, 0, 3), + cute.group_modes(tSgQ_qdl, 0, 3), + ) + kv_cta_layout = cute.make_layout( + cute.slice_(cluster_layout_vmnk, (0, None, 0, 0)).shape + ) + # (bN, bK, loopN, loopK, loopL) + gK_kdl = cute.flat_divide(mK_kdl, cute.select(self.qk_mma_tiler, mode=[1, 2])) + tSgK_kdl = qk_thr_mma.partition_B(gK_kdl) + tKsK, tKgK_kdl = cute.nvgpu.cpasync.tma_partition( + tma_atom_k, + block_in_cluster_coord_vmnk[1], + kv_cta_layout, + cute.group_modes(sK, 0, 3), + cute.group_modes(tSgK_kdl, 0, 3), + ) + # (blk, loopBlk, loopL) + gScaleK_kdl = cute.flat_divide(mScaleK_kdl, self.scale_k_tiler) + # Deal with 2cta + gScaleK_kdl_ = cute.logical_divide(gScaleK_kdl, (self.scale_k_tiler[0] // 2,))[ + (None, mma_tile_coord_v), None, None + ] + tKsScaleK, tKgScaleK_kdl = cute.nvgpu.cpasync.tma_partition( + tma_atom_scale_k, + block_in_cluster_coord_vmnk[1], + kv_cta_layout, + sScaleK, + gScaleK_kdl_, + ) + # (bN, bK, loopN, loopK, loopL) + gV_dkl = cute.flat_divide(mV_dkl, cute.select(self.pv_mma_tiler, mode=[1, 2])) + tOgV_dkl = pv_thr_mma.partition_B(gV_dkl) + tVsV, tVgV_dkl = cute.nvgpu.cpasync.tma_partition( + tma_atom_v, + block_in_cluster_coord_vmnk[1], + kv_cta_layout, + cute.group_modes(sV, 0, 3), + cute.group_modes(tOgV_dkl, 0, 3), + ) + # (bBlk, loopBlk, loopL) + gScaleV_dkl = cute.flat_divide(mScaleV_dkl, self.scale_v_tiler) + tVsScaleV, tVgScaleV_dkl = cute.nvgpu.cpasync.tma_partition( + tma_atom_scale_v, + block_in_cluster_coord_vmnk[1], + kv_cta_layout, + sScaleV, + gScaleV_dkl, + ) + # (bM, bN, loopM, loopN, loopL) + gO_qdl = cute.flat_divide(mO_qdl, cute.select(self.pv_block_tiler, mode=[0, 1])) + cO_qdl = cute.flat_divide( + cute.make_identity_tensor(mO_qdl.shape), + cute.select(self.pv_block_tiler, mode=[0, 1]), + ) + seqlen_q = mQ_qdl.shape[0] + seqlen_k = mK_kdl.shape[0] + tile_sched = fmha_utils.create_fmha_static_tile_scheduler( + tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim() + ) + work_tile = tile_sched.initial_work_tile_info() + # Cluster wait + pipeline_init_wait(cluster_shape_mn=cluster_layout_vmnk) + + # /////////////////////////////////////////////////////////////////////////////// + # Load + # /////////////////////////////////////////////////////////////////////////////// + if warp_idx == self.load_warp_id: + cute.arch.setmaxregister_decrease(self.num_regs_other) + while work_tile.is_valid_tile: + curr_block_coord = work_tile.tile_idx + mma_block_coord = ( + curr_block_coord[0] // cute.size(qk_tiled_mma.thr_id.shape), + curr_block_coord[1], + curr_block_coord[2], + ) + seqlen_kv_loop_steps = fmha_utils.FusedMask.get_trip_count( + self.mask_type, + mma_block_coord, + self.qk_mma_tiler, + seqlen_q, + seqlen_k, + window_size_left, + window_size_right, + ) + # ((atom_v, rest_v), RestK) + tQgQ = tQgQ_qdl[None, mma_block_coord[0], None, mma_block_coord[2]] + # ((atom_v, rest_v), RestN, RestK) + tKgK = tKgK_kdl[None, None, None, mma_block_coord[2]] + tKgScaleK = tKgScaleK_kdl[None, None, mma_block_coord[2]] + # ((atom_v, rest_v), RestN, RestK) + tVgV = tVgV_dkl[None, None, None, mma_block_coord[2]] + tVgScaleV = tVgScaleV_dkl[None, None, mma_block_coord[2]] + load_kv_producer, load_scale_k_producer, load_q_producer = ( + prefill_utils.load_qk( # Q & K0 & ScaleK0 + iterations=self.iterations_qk, + kv_step=0, + k_args=(tKgK, tKsK, tma_atom_k, load_kv_producer), + scale_k_args=( + tKgScaleK, + tKsScaleK, + tma_atom_scale_k, + load_scale_k_producer, + ), + q_args=(tQgQ, tQsQ, tma_atom_q, load_q_producer), + ) + ) + for step in cutlass.range(1, seqlen_kv_loop_steps, 1, unroll=1): + load_kv_producer, load_scale_k_producer = ( + prefill_utils.load_qk( # Ki & ScaleKi + iterations=self.iterations_qk, + kv_step=step, + k_args=(tKgK, tKsK, tma_atom_k, load_kv_producer), + scale_k_args=( + tKgScaleK, + tKsScaleK, + tma_atom_scale_k, + load_scale_k_producer, + ), + ) + ) + load_kv_producer, load_scale_v_producer = ( + prefill_utils.load_v( # Vi-1 & ScaleVi-1 + iterations=self.iterations_pv, + kv_step=step - 1, + v_args=(tVgV, tVsV, tma_atom_v, load_kv_producer), + scale_v_args=( + tVgScaleV, + tVsScaleV, + tma_atom_scale_v, + load_scale_v_producer, + ), + ) + ) + load_kv_producer, load_scale_v_producer = ( + prefill_utils.load_v( # Vend & ScaleVend + iterations=self.iterations_pv, + kv_step=seqlen_kv_loop_steps - 1, + v_args=(tVgV, tVsV, tma_atom_v, load_kv_producer), + scale_v_args=( + tVgScaleV, + tVsScaleV, + tma_atom_scale_v, + load_scale_v_producer, + ), + ) + ) + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + load_kv_producer.tail() + load_scale_k_producer.tail() + load_scale_v_producer.tail() + load_q_producer.tail() + + # /////////////////////////////////////////////////////////////////////////////// + # MMA + # /////////////////////////////////////////////////////////////////////////////// + if warp_idx == self.mma_warp_id: + cute.arch.setmaxregister_decrease(self.num_regs_other) + tmem.wait_for_alloc() + while work_tile.is_valid_tile: + curr_block_coord = work_tile.tile_idx + mma_block_coord = ( + curr_block_coord[0] // cute.size(qk_tiled_mma.thr_id.shape), + curr_block_coord[1], + curr_block_coord[2], + ) + seqlen_kv_loop_steps = fmha_utils.FusedMask.get_trip_count( + self.mask_type, + mma_block_coord, + self.qk_mma_tiler, + seqlen_q, + seqlen_k, + window_size_left, + window_size_right, + ) + load_q_releaser = load_q_consumer.clone() + pv_tiled_mma.set(tcgen05.Field.ACCUMULATE, False) + if seqlen_kv_loop_steps > 1: + mma_s_producer, load_q_consumer, dequant_kv_consumer = ( + prefill_utils.mma_qk( # QK0 + self.iterations_qk, + qk_tiled_mma, + (tStS, tSrQ, tSrK_trans), + ( + mma_s_producer, + load_q_consumer, + None, + dequant_kv_consumer, + ), + ) + ) + for i in cutlass.range(1, seqlen_kv_loop_steps - 1, 1, unroll=1): + mma_s_producer, _, dequant_kv_consumer = ( + prefill_utils.mma_qk( # QKi + self.iterations_qk, + qk_tiled_mma, + (tStS, tSrQ, tSrK_trans), + (mma_s_producer, None, None, dequant_kv_consumer), + ) + ) + ( + pv_tiled_mma, + p_mma_consumer, + mma_o_producer, + dequant_kv_consumer, + swap_producer, + ) = self.mma_pv( # PVi + pv_tiled_mma, + (tOtO_staged, tOrP, tOrV_trans), + ( + p_mma_consumer, + mma_o_producer, + dequant_kv_consumer, + swap_producer, + ), + ) + mma_s_producer, _, dequant_kv_consumer = ( + prefill_utils.mma_qk( # QKend needs to release Q + self.iterations_qk, + qk_tiled_mma, + (tStS, tSrQ, tSrK_trans), + ( + mma_s_producer, + None, + load_q_releaser, + dequant_kv_consumer, + ), + ) + ) + ( + pv_tiled_mma, + p_mma_consumer, + mma_o_producer, + dequant_kv_consumer, + swap_producer, + ) = self.mma_pv( # PVend-1 + pv_tiled_mma, + (tOtO_staged, tOrP, tOrV_trans), + ( + p_mma_consumer, + mma_o_producer, + dequant_kv_consumer, + swap_producer, + ), + ) + else: + mma_s_producer, load_q_consumer, dequant_kv_consumer = ( + prefill_utils.mma_qk( # QK0 + self.iterations_qk, + qk_tiled_mma, + (tStS, tSrQ, tSrK_trans), + ( + mma_s_producer, + load_q_consumer, + load_q_releaser, + dequant_kv_consumer, + ), + ) + ) + ( + pv_tiled_mma, + p_mma_consumer, + mma_o_producer, + dequant_kv_consumer, + swap_producer, + ) = self.mma_pv( # PVend + pv_tiled_mma, + (tOtO_staged, tOrP, tOrV_trans), + ( + p_mma_consumer, + mma_o_producer, + dequant_kv_consumer, + swap_producer, + ), + ) + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + mma_s_producer.tail() + mma_o_producer.tail() + swap_producer.tail() + + # /////////////////////////////////////////////////////////////////////////////// + # Softmax + # /////////////////////////////////////////////////////////////////////////////// + if warp_idx < self.mma_warp_id and warp_idx >= self.softmax_warp_ids[0]: + cute.arch.setmaxregister_increase(self.num_regs_softmax) + tmem.allocate(self.num_tmem_alloc_cols) + tmem.wait_for_alloc() + tmem_ptr = tmem.retrieve_ptr(self.qk_acc_dtype) + while work_tile.is_valid_tile: + curr_block_coord = work_tile.tile_idx + mma_block_coord = ( + curr_block_coord[0] // cute.size(qk_tiled_mma.thr_id.shape), + curr_block_coord[1], + curr_block_coord[2], + ) + seqlen_kv_loop_steps = fmha_utils.FusedMask.get_trip_count( + self.mask_type, + mma_block_coord, + self.qk_mma_tiler, + seqlen_q, + seqlen_k, + window_size_left, + window_size_right, + ) + unmask_steps = fmha_utils.FusedMask.get_unmasked_trip_count( + self.mask_type, + mma_block_coord, + self.qk_mma_tiler, + seqlen_q, + seqlen_k, + window_size_left, + window_size_right, + ) + gO_staged = gO_qdl[ + None, None, curr_block_coord[0], None, curr_block_coord[2] + ] + cO_staged = cO_qdl[ + None, None, curr_block_coord[0], None, curr_block_coord[2] + ] + cS_base = cute.make_identity_tensor( + (self.qk_mma_tiler[0], self.qk_mma_tiler[1]) + ) + cS = cute.domain_offset( + (mma_block_coord[0] * self.qk_mma_tiler[0], 0), cS_base + ) + tScS = qk_thr_mma.partition_C(cS) + row_max = -Float32.inf + row_max_prev = -Float32.inf + row_sum = 0.0 + # S0 -> P0 + row_max, row_sum, mma_s_consumer, p_mma_producer = self.softmax_step( + (unmask_steps == 0, window_size_left, window_size_right), + (row_max, row_sum, seqlen_q, seqlen_k, scale_softmax_log2), + (tStS, tScS, sP), + (mma_s_consumer, p_mma_producer), + ) + row_max_prev = row_max + # Use stage1 to do S/O swap + tmem_tiled_load, tmem_tiled_store, tSWAPtO, tSWAPrO = ( + self.get_swap_o_partition(tOtO_staged, cO_staged) + ) + for step in cutlass.range(1, seqlen_kv_loop_steps, 1, unroll=1): + cS_iter = cute.domain_offset((0, step * self.qk_mma_tiler[1]), cS) + tScS_iter = qk_thr_mma.partition_C(cS_iter) + # Si -> Pi & Oi-1 -> Oi + ( + tSWAPrO, + row_max, + row_sum, + mma_s_consumer, + p_mma_producer, + mma_o_consumer, + swap_consumer, + ) = self.softmax_correction_step( + ( + step >= unmask_steps, + step > 1, + step == seqlen_kv_loop_steps - 1, + window_size_left, + window_size_right, + ), + ( + row_max_prev, + row_sum, + seqlen_q, + seqlen_k, + scale_softmax_log2, + ), + (tStS, tScS_iter, sP, tOtO_staged, cO_staged), + (mma_s_consumer, p_mma_producer, mma_o_consumer, swap_consumer), + (tmem_tiled_load, tmem_tiled_store, tSWAPtO, tSWAPrO), + epi_tile, + ) + row_max_prev = row_max + # O_partial -> O_final + mma_o_consumer, swap_consumer = self.correction_epilog( + (tmem_tiled_store, tSWAPrO, tSWAPtO, swap_consumer), + (row_sum, seqlen_q, scale_output), + (mma_o_consumer, gO_staged, cO_staged, tOtO_staged), + epi_tile, + ) + # Make sure we start the next wave's S=QK after correction epilog + mma_s_consumer.release() + mma_s_consumer.advance() + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + p_mma_producer.tail() + tmem.relinquish_alloc_permit() + tmem.free(tmem_ptr) + + # /////////////////////////////////////////////////////////////////////////////// + # Trans + # /////////////////////////////////////////////////////////////////////////////// + if warp_idx < self.softmax_warp_ids[0]: + cute.arch.setmaxregister_decrease(self.num_regs_transform) + qk_thr_mma_leader_cta = qk_tiled_mma.get_slice(0) + pv_thr_mma_leader_cta = pv_tiled_mma.get_slice(0) + sScaleK_ = qk_thr_mma_leader_cta.partition_B(sScaleK_s2r_view) + sScaleV_ = pv_thr_mma_leader_cta.partition_B(sScaleV_s2r_view) + while work_tile.is_valid_tile: + curr_block_coord = work_tile.tile_idx + mma_block_coord = ( + curr_block_coord[0] // cute.size(qk_tiled_mma.thr_id.shape), + curr_block_coord[1], + curr_block_coord[2], + ) + seqlen_kv_loop_steps = fmha_utils.FusedMask.get_trip_count( + self.mask_type, + mma_block_coord, + self.qk_mma_tiler, + seqlen_q, + seqlen_k, + window_size_left, + window_size_right, + ) + load_kv_consumer, load_scale_k_consumer, dequant_kv_producer = ( + prefill_utils.dequant_k( # K0 + self.iterations_qk, + self.transform_warp_ids, + (self.k_dtype, self.q_dtype), + (sK, sScaleK_, sK_trans), + (load_kv_consumer, load_scale_k_consumer, dequant_kv_producer), + ) + ) + for step in cutlass.range(1, seqlen_kv_loop_steps, 1, unroll=1): + load_kv_consumer, load_scale_k_consumer, dequant_kv_producer = ( + prefill_utils.dequant_k( # Ki + self.iterations_qk, + self.transform_warp_ids, + (self.k_dtype, self.q_dtype), + (sK, sScaleK_, sK_trans), + ( + load_kv_consumer, + load_scale_k_consumer, + dequant_kv_producer, + ), + ) + ) + load_kv_consumer, load_scale_v_consumer, dequant_kv_producer = ( + prefill_utils.dequant_v( # Vi-1 + self.iterations_pv, + self.transform_warp_ids, + (self.v_dtype, self.q_dtype), + (sV, sScaleV_, sV_trans), + ( + load_kv_consumer, + load_scale_v_consumer, + dequant_kv_producer, + ), + ) + ) + load_kv_consumer, load_scale_v_consumer, dequant_kv_producer = ( + prefill_utils.dequant_v( # Vend + self.iterations_pv, + self.transform_warp_ids, + (self.v_dtype, self.q_dtype), + (sV, sScaleV_, sV_trans), + (load_kv_consumer, load_scale_v_consumer, dequant_kv_producer), + ) + ) + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + dequant_kv_producer.tail() + + # /////////////////////////////////////////////////////////////////////////////// + # Empty + # /////////////////////////////////////////////////////////////////////////////// + if warp_idx > self.load_warp_id: + cute.arch.setmaxregister_decrease(self.num_regs_other) + + return + + @cute.jit + def get_swap_o_partition( + self, + tOtO_staged: cute.Tensor, + cO_staged: cute.Tensor, + ) -> cute.Tensor: + # Swap S & O on stage 1 + tOtO_stage1 = tOtO_staged[(None, None), 0, 0, 1] + cO_stage1 = cO_staged[None, None, 1] + tmem_load_atom = cute.make_copy_atom( + tcgen05.Ld32x32bOp(tcgen05.Repetition(32)), self.pv_acc_dtype + ) + tmem_tiled_load = tcgen05.make_tmem_copy(tmem_load_atom, tOtO_stage1) + tidx, _, _ = cute.arch.thread_idx() + thread_idx = tidx % (self.threads_per_warp * len(self.softmax_warp_ids)) + thr_load = tmem_tiled_load.get_slice(thread_idx) + tSWAPtO = thr_load.partition_D(tOtO_stage1) + tSWAPcO = thr_load.partition_D(cO_stage1) + tSWAPrO = cute.make_rmem_tensor(tSWAPcO.shape, self.pv_acc_dtype) + tmem_store_atom = cute.make_copy_atom( + tcgen05.St32x32bOp(tcgen05.Repetition(32)), self.pv_acc_dtype + ) + tmem_tiled_store = tcgen05.make_tmem_copy(tmem_store_atom, tOtO_staged) + return tmem_tiled_load, tmem_tiled_store, tSWAPtO, tSWAPrO + + @cute.jit + def mma_pv( + self, + pv_tiled_mma: cute.TiledMma, + tensor_args: Tuple, + pipeline_args: Tuple, + ): + tOtO_staged, tOrP, tOrV_trans = tensor_args + p_mma_consumer, mma_o_producer, dequant_kv_consumer, swap_producer = ( + pipeline_args + ) + cta_rank_in_cluster = cute.arch.make_warp_uniform( + cute.arch.block_idx_in_cluster() + ) + is_leader_cta = cta_rank_in_cluster % 2 == 0 + if is_leader_cta: + p_handle = p_mma_consumer.wait_and_advance() + o_handle = mma_o_producer.acquire_and_advance() + pv_whether_acc = pv_tiled_mma.get(tcgen05.Field.ACCUMULATE) + for iter_n in cutlass.range(self.iterations_pv, unroll=1): + v_trans_handle = dequant_kv_consumer.wait_and_advance() + if iter_n == 1: + swap_producer.acquire() + pv_tiled_mma.set(tcgen05.Field.ACCUMULATE, pv_whether_acc) + tOtO_slice = tOtO_staged[None, None, None, iter_n] + tOrP_slice = tOrP[None, None, None, p_handle.index] + tOrV_trans_slice = tOrV_trans[None, None, None, v_trans_handle.index] + num_kphases = cute.size(tOrV_trans_slice, mode=[2]) + for kphase_idx in cutlass.range(num_kphases, unroll_full=True): + kphase_coord = (None, None, kphase_idx) + cute.gemm( + pv_tiled_mma, + tOtO_slice, + tOrP_slice[kphase_coord], + tOrV_trans_slice[kphase_coord], + tOtO_slice, + ) + pv_tiled_mma.set(tcgen05.Field.ACCUMULATE, True) + if iter_n == 1: + swap_producer.commit() + swap_producer.advance() + v_trans_handle.release() + o_handle.commit() + p_handle.release() + return ( + pv_tiled_mma, + p_mma_consumer, + mma_o_producer, + dequant_kv_consumer, + swap_producer, + ) + + @cute.jit + def softmax_step( + self, + mask_args: Tuple, + value_args: Tuple, + tensor_args: Tuple, + pipeline_args: Tuple, + ) -> Tuple[Float32, Float32, pipeline.PipelineConsumer, pipeline.PipelineProducer]: + need_apply_mask, window_size_left, window_size_right = mask_args + row_max, row_sum, seqlen_q, seqlen_k, scale_softmax_log2 = value_args + tStS, tScS, sP = tensor_args + mma_s_consumer, p_mma_producer = pipeline_args + tidx, _, _ = cute.arch.thread_idx() + thread_idx = tidx % (self.threads_per_warp * len(self.softmax_warp_ids)) + s_handle = mma_s_consumer.wait_and_advance() + tStS_slice = tStS[(None, None), 0, 0, s_handle.index] + tScS_slice = tScS[(None, None), 0, 0] + tmem_load_atom = cute.make_copy_atom( + tcgen05.Ld32x32bOp(tcgen05.Repetition(32)), self.qk_acc_dtype + ) + tmem_tiled_load = tcgen05.make_tmem_copy(tmem_load_atom, tStS_slice) + thr_load = tmem_tiled_load.get_slice(thread_idx) + tTMEM_LOADtS = thr_load.partition_S(tStS_slice) + tTMEM_LOADcS = thr_load.partition_D(tScS_slice) + tTMEM_LOADrS = cute.make_rmem_tensor(tTMEM_LOADcS.shape, self.qk_acc_dtype) + cute.copy(tmem_tiled_load, tTMEM_LOADtS, tTMEM_LOADrS) + cute.arch.fence_view_async_tmem_load() + s_handle.release() + if need_apply_mask: + fmha_utils.FusedMask.apply_mask( + self.mask_type, + tTMEM_LOADrS, + tTMEM_LOADcS, + seqlen_q, + seqlen_k, + window_size_left, + window_size_right, + ) + old_row_max = row_max + row_max = tTMEM_LOADrS.load().reduce(cute.ReductionOp.MAX, row_max, 0) + row_max_safe = row_max + if row_max == -cutlass.Float32.inf: + row_max_safe = 0.0 + scale = scale_softmax_log2 + minus_row_max_scale = (0.0 - row_max_safe) * scale + tTMEM_STORErP = cute.make_rmem_tensor(tTMEM_LOADrS.shape, self.p_dtype) + for k in cutlass.range(cute.size(tTMEM_LOADrS), vectorize=True): + tTMEM_LOADrS[k] = tTMEM_LOADrS[k] * scale + minus_row_max_scale + tTMEM_LOADrS[k] = cute.math.exp2(tTMEM_LOADrS[k], fastmath=True) + s_vec = tTMEM_LOADrS.load() + tTMEM_STORErP.store(s_vec.to(self.p_dtype)) + + p_handle = p_mma_producer.acquire_and_advance() + sP_slice = sP[None, None, None, p_handle.index] + sP_mk_view = cute.make_tensor( + sP_slice.iterator, + cute.make_layout( + ( + (sP_slice.shape[0][0], sP_slice.shape[1]), + (sP_slice.shape[0][1], sP_slice.shape[2]), + ), + stride=( + (sP_slice.stride[0][0], sP_slice.stride[1]), + (sP_slice.stride[0][1], sP_slice.stride[2]), + ), + ), + ) + universal_copy_bits = 128 + smem_copy_atom = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), + self.q_dtype, + num_bits_per_copy=universal_copy_bits, + ) + smem_tiled_copy = cute.make_tiled_copy_D(smem_copy_atom, tmem_tiled_load) + smem_thr_copy = smem_tiled_copy.get_slice(thread_idx) + rP_copy_view = smem_thr_copy.retile(tTMEM_STORErP) + sP_copy_view = smem_thr_copy.partition_D(sP_mk_view) + cute.copy(smem_tiled_copy, rP_copy_view, sP_copy_view) + cute.arch.fence_view_async_shared() + p_handle.commit() + acc_scale_ = scale * (old_row_max - row_max_safe) + acc_scale = cute.math.exp2(acc_scale_, fastmath=True) * 0.5 + # TODO: calc row sum with TensorSSA + row_sum *= acc_scale + local_row_sum_0 = (row_sum, row_sum) + local_row_sum_1 = (0.0, 0.0) + local_row_sum_2 = (0.0, 0.0) + local_row_sum_3 = (0.0, 0.0) + reduction_unroll = 4 + frg_tile = cute.size(tTMEM_LOADrS) // reduction_unroll + tTMEM_LOADrS_frg = cute.logical_divide(tTMEM_LOADrS, cute.make_layout(frg_tile)) + for j in cutlass.range_constexpr(0, cute.size(tTMEM_LOADrS_frg, mode=[0]), 2): + local_row_sum_0 = cute.arch.add_packed_f32x2( + local_row_sum_0, (tTMEM_LOADrS_frg[j, 0], tTMEM_LOADrS_frg[j + 1, 0]) + ) + local_row_sum_1 = cute.arch.add_packed_f32x2( + local_row_sum_1, (tTMEM_LOADrS_frg[j, 1], tTMEM_LOADrS_frg[j + 1, 1]) + ) + local_row_sum_2 = cute.arch.add_packed_f32x2( + local_row_sum_2, (tTMEM_LOADrS_frg[j, 2], tTMEM_LOADrS_frg[j + 1, 2]) + ) + local_row_sum_3 = cute.arch.add_packed_f32x2( + local_row_sum_3, (tTMEM_LOADrS_frg[j, 3], tTMEM_LOADrS_frg[j + 1, 3]) + ) + local_row_sum_0 = cute.arch.add_packed_f32x2(local_row_sum_0, local_row_sum_1) + local_row_sum_2 = cute.arch.add_packed_f32x2(local_row_sum_2, local_row_sum_3) + local_row_sum_0 = cute.arch.add_packed_f32x2(local_row_sum_0, local_row_sum_2) + row_sum = local_row_sum_0[0] + local_row_sum_0[1] + return row_max, row_sum, mma_s_consumer, p_mma_producer + + @cute.jit + def correction_rescale( + self, + tOtO: cute.Tensor, + cO: cute.Tensor, + epi_tile: cute.Tile, + scale: Float32, + ) -> cute.Tensor: + tidx, _, _ = cute.arch.thread_idx() + thread_idx = tidx % (self.threads_per_warp * len(self.softmax_warp_ids)) + tOtO_epi = cute.zipped_divide(tOtO, cute.make_layout(epi_tile)) + cO_epi = cute.zipped_divide(cO, cute.make_layout(epi_tile)) + tmem_load_atom = cute.make_copy_atom( + tcgen05.Ld32x32bOp(tcgen05.Repetition(32)), + self.pv_acc_dtype, + ) + tmem_tiled_load = tcgen05.make_tmem_copy(tmem_load_atom, tOtO_epi) + thr_load = tmem_tiled_load.get_slice(thread_idx) + tmem_store_atom = cute.make_copy_atom( + tcgen05.St32x32bOp(tcgen05.Repetition(32)), + self.pv_acc_dtype, + ) + tmem_store_atom = tcgen05.make_tmem_copy(tmem_store_atom, tOtO_epi) + thr_store = tmem_store_atom.get_slice(thread_idx) + tTMEM_LOADtO = thr_load.partition_S(tOtO_epi) + tTMEM_LOADcO = thr_load.partition_D(cO_epi) + tTMEM_STOREtO = thr_store.partition_D(tOtO_epi) + + tTMrO = cute.make_rmem_tensor_like( + cute.append( + cute.make_layout(tTMEM_LOADcO[None, 0, 0].shape), + cute.make_layout(2, stride=cute.size(tTMEM_LOADcO[None, 0, 0].shape)), + ), + self.pv_acc_dtype, + ) + tTMEM_LOADtO_0 = tTMEM_LOADtO[None, 0, 0] + cute.copy(tmem_tiled_load, tTMEM_LOADtO_0, tTMrO[None, 0]) + iter_num = cute.size(tTMEM_LOADtO, mode=[1]) + for i in cutlass.range(1, iter_num, unroll_full=True): + tTMEM_LOADtO_i = tTMEM_LOADtO[None, i, 0] + cute.copy(tmem_tiled_load, tTMEM_LOADtO_i, tTMrO[None, i % 2]) + for j in cutlass.range( + cute.size(tTMrO, mode=[0]), unroll_full=True, vectorize=True + ): + tTMrO[j, (i - 1) % 2] = tTMrO[j, (i - 1) % 2] * scale + tTMEM_STOREtO_prev_i = tTMEM_STOREtO[None, i - 1, 0] + cute.copy(tmem_store_atom, tTMrO[None, (i - 1) % 2], tTMEM_STOREtO_prev_i) + for j in cutlass.range( + cute.size(tTMrO, mode=[0]), unroll_full=True, vectorize=True + ): + tTMrO[j, (iter_num - 1) % 2] = tTMrO[j, (iter_num - 1) % 2] * scale + cute.copy( + tmem_store_atom, + tTMrO[None, (iter_num - 1) % 2], + tTMEM_STOREtO[None, iter_num - 1, 0], + ) + + @cute.jit + def sum_reduction( + self, + tensor: cute.Tensor, + ) -> Float32: + local_sum_0 = (0.0, 0.0) + local_sum_1 = (0.0, 0.0) + local_sum_2 = (0.0, 0.0) + local_sum_3 = (0.0, 0.0) + reduction_unroll = 4 + frg_tile = cute.size(tensor) // reduction_unroll + tensor_frg = cute.logical_divide(tensor, cute.make_layout(frg_tile)) + for i in cutlass.range(0, cute.size(tensor_frg, mode=[0]), 2, unroll_full=True): + local_sum_0 = cute.arch.add_packed_f32x2( + local_sum_0, (tensor_frg[i, 0], tensor_frg[i + 1, 0]) + ) + local_sum_1 = cute.arch.add_packed_f32x2( + local_sum_1, (tensor_frg[i, 1], tensor_frg[i + 1, 1]) + ) + local_sum_2 = cute.arch.add_packed_f32x2( + local_sum_2, (tensor_frg[i, 2], tensor_frg[i + 1, 2]) + ) + local_sum_3 = cute.arch.add_packed_f32x2( + local_sum_3, (tensor_frg[i, 3], tensor_frg[i + 1, 3]) + ) + local_sum_0 = cute.arch.add_packed_f32x2(local_sum_0, local_sum_1) + local_sum_2 = cute.arch.add_packed_f32x2(local_sum_2, local_sum_3) + local_sum_0 = cute.arch.add_packed_f32x2(local_sum_0, local_sum_2) + return local_sum_0[0] + local_sum_0[1] + + @cute.jit + def softmax_correction_step( + self, + mask_args: Tuple, + value_args: Tuple, + tensor_args: Tuple, + pipeline_args: Tuple, + swap_args: Tuple, + epi_tile: cute.Tile, + ) -> Tuple[Float32, Float32, pipeline.PipelineConsumer, pipeline.PipelineProducer]: + ( + need_apply_mask, + need_store_o, + last_iteration, + window_size_left, + window_size_right, + ) = mask_args + row_max, row_sum, seqlen_q, seqlen_k, scale_softmax_log2 = value_args + tStS, tScS_iter, sP, tOtO_staged, cO_staged = tensor_args + mma_s_consumer, p_mma_producer, mma_o_consumer, swap_consumer = pipeline_args + swap_tmem_tiled_load, swap_tmem_tiled_store, tSWAPtO, tSWAPrO = swap_args + + tidx, _, _ = cute.arch.thread_idx() + thread_idx = tidx % (self.threads_per_warp * len(self.softmax_warp_ids)) + tStS_slice = tStS[(None, None), 0, 0, None] + tScS_slice = tScS_iter[(None, None), 0, 0] + tmem_load_atom = cute.make_copy_atom( + tcgen05.Ld32x32bOp(tcgen05.Repetition(32)), self.qk_acc_dtype + ) + tmem_tiled_load = tcgen05.make_tmem_copy( + tmem_load_atom, tStS_slice[None, None, 0] + ) + thr_load = tmem_tiled_load.get_slice(thread_idx) + tTMEM_LOADtS = thr_load.partition_S(tStS_slice) + tTMEM_LOADcS = thr_load.partition_D(tScS_slice) + tTMEM_LOADrS = cute.make_rmem_tensor(tTMEM_LOADcS.shape, self.qk_acc_dtype) + mma_s_consumer.wait() + if not need_store_o: + cute.copy( + tmem_tiled_load, + tTMEM_LOADtS[None, None, None, 0], + tTMEM_LOADrS, + ) + else: + iter_num = cute.size(tTMEM_LOADrS, mode=[2]) + for i in cutlass.range(iter_num, unroll_full=True): + cute.copy( + swap_tmem_tiled_load, + tTMEM_LOADtS[None, 0, iter_num - 1 - i, 0], + tTMEM_LOADrS[None, 0, iter_num - 1 - i], + ) + cute.copy( + swap_tmem_tiled_store, + tSWAPrO[None, 0, iter_num - 1 - i], + tSWAPtO[None, 0, iter_num - 1 - i], + ) + cute.arch.fence_view_async_tmem_load() + cute.arch.fence_view_async_tmem_store() + swap_consumer.release() + swap_consumer.advance() + + if need_apply_mask: + fmha_utils.FusedMask.apply_mask( + self.mask_type, + tTMEM_LOADrS, + tTMEM_LOADcS, + seqlen_q, + seqlen_k, + window_size_left, + window_size_right, + ) + old_row_max = row_max + row_max = tTMEM_LOADrS.load().reduce(cute.ReductionOp.MAX, row_max, 0) + row_max_safe = row_max + if row_max == -cutlass.Float32.inf: + row_max_safe = 0.0 + scale = scale_softmax_log2 + minus_row_max_scale = (0.0 - row_max_safe) * scale + acc_scale_ = scale * (old_row_max - row_max_safe) + acc_scale = cute.math.exp2(acc_scale_, fastmath=True) + row_sum *= acc_scale + + tTMEM_STORErP = cute.make_rmem_tensor(tTMEM_LOADrS.shape, self.p_dtype) + subtile_cnt = 2 + subtile_size = cute.size(tTMEM_LOADrS) // subtile_cnt + tTMEM_LOADrS_subtile = cute.logical_divide( + tTMEM_LOADrS, cute.make_layout(subtile_size) + ) + tTMEM_STORErP_subtile = cute.logical_divide( + tTMEM_STORErP, cute.make_layout(subtile_size) + ) + + tTMEM_LOADrS_0 = tTMEM_LOADrS_subtile[None, 0] + tTMEM_STORErP_0 = tTMEM_STORErP_subtile[None, 0] + for k in cutlass.range(cute.size(tTMEM_LOADrS_0), vectorize=True): + tTMEM_LOADrS_0[k] = tTMEM_LOADrS_0[k] * scale + minus_row_max_scale + tTMEM_LOADrS_0[k] = cute.math.exp2(tTMEM_LOADrS_0[k], fastmath=True) + row_sum += self.sum_reduction(tTMEM_LOADrS_0) + s_vec = tTMEM_LOADrS_0.load() + tTMEM_STORErP_0.store(s_vec.to(self.p_dtype)) + + swap_consumer.wait() + cute.copy(swap_tmem_tiled_load, tSWAPtO, tSWAPrO) + cute.arch.fence_view_async_tmem_load() + if not last_iteration: + mma_s_consumer.release() + mma_s_consumer.advance() + + tTMEM_LOADrS_1 = tTMEM_LOADrS_subtile[None, 1] + tTMEM_STORErP_1 = tTMEM_STORErP_subtile[None, 1] + for k in cutlass.range(cute.size(tTMEM_LOADrS_1), vectorize=True): + tTMEM_LOADrS_1[k] = tTMEM_LOADrS_1[k] * scale + minus_row_max_scale + tTMEM_LOADrS_1[k] = cute.math.exp2(tTMEM_LOADrS_1[k], fastmath=True) + row_sum += self.sum_reduction(tTMEM_LOADrS_1) + s_vec = tTMEM_LOADrS_1.load() + tTMEM_STORErP_1.store(s_vec.to(self.p_dtype)) + + p_handle = p_mma_producer.acquire_and_advance() + sP_slice = sP[None, None, None, p_handle.index] + sP_mk_view = cute.make_tensor( + sP_slice.iterator, + cute.make_layout( + ( + (sP_slice.shape[0][0], sP_slice.shape[1]), + (sP_slice.shape[0][1], sP_slice.shape[2]), + ), + stride=( + (sP_slice.stride[0][0], sP_slice.stride[1]), + (sP_slice.stride[0][1], sP_slice.stride[2]), + ), + ), + ) + universal_copy_bits = 128 + smem_copy_atom = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), + self.q_dtype, + num_bits_per_copy=universal_copy_bits, + ) + smem_tiled_copy = cute.make_tiled_copy_D(smem_copy_atom, tmem_tiled_load) + smem_thr_copy = smem_tiled_copy.get_slice(thread_idx) + rP_copy_view = smem_thr_copy.retile(tTMEM_STORErP) + sP_copy_view = smem_thr_copy.partition_D(sP_mk_view) + cute.copy(smem_tiled_copy, rP_copy_view, sP_copy_view) + for i in cutlass.range(0, cute.size(tSWAPrO), 2, unroll_full=True): + tSWAPrO[i], tSWAPrO[i + 1] = cute.arch.mul_packed_f32x2( + (tSWAPrO[i], tSWAPrO[i + 1]), + (acc_scale, acc_scale), + ) + cute.arch.fence_view_async_shared() + p_handle.commit() + o_handle = mma_o_consumer.wait_and_advance() + self.correction_rescale( + tOtO_staged[(None, None), 0, 0, 0], + cO_staged[None, None, 0], + epi_tile, + acc_scale, + ) + # Skip stage1; loop starts at 2 + for iter_n in cutlass.range(2, self.iterations_pv, unroll_full=True): + self.correction_rescale( + tOtO_staged[(None, None), 0, 0, iter_n], + cO_staged[None, None, iter_n], + epi_tile, + acc_scale, + ) + cute.arch.fence_view_async_tmem_store() + o_handle.release() + + return ( + tSWAPrO, + row_max, + row_sum, + mma_s_consumer, + p_mma_producer, + mma_o_consumer, + swap_consumer, + ) + + @cute.jit + def correction_epilog( + self, + swap_args: Tuple, + value_args: Tuple, + o_args: Tuple, + epi_tile: cute.Tile, + ) -> Tuple[pipeline.PipelineConsumer, pipeline.PipelineProducer]: + (swap_tmem_tiled_store, tSWAPrO, tSWAPtO, swap_consumer) = swap_args + (row_sum, seqlen_q, scale_output) = value_args + (mma_o_consumer, gO_staged, cO_staged, tOtO_staged) = o_args + + cute.copy(swap_tmem_tiled_store, tSWAPrO, tSWAPtO) + cute.arch.fence_view_async_tmem_store() + swap_consumer.release() + swap_consumer.advance() + tidx, _, _ = cute.arch.thread_idx() + thread_idx = tidx % (self.threads_per_warp * len(self.softmax_warp_ids)) + scale = scale_output / row_sum + o_handle = mma_o_consumer.wait_and_advance() + # empty step as we access tOtO_stage1 by the normal way + swap_consumer.wait() + swap_consumer.release() + swap_consumer.advance() + for iter_n in cutlass.range(self.iterations_pv): + gO = gO_staged[None, None, iter_n] + cO = cO_staged[None, None, iter_n] + tOtO = tOtO_staged[(None, None), 0, 0, iter_n] + tOtO_epi = cute.zipped_divide(tOtO, epi_tile) + cO_epi = cute.zipped_divide(cO, epi_tile) + gO_epi = cute.zipped_divide(gO, epi_tile) + tidx, _, _ = cute.arch.thread_idx() + thread_idx = tidx % (self.threads_per_warp * len(self.softmax_warp_ids)) + tmem_copy_atom = cute.make_copy_atom( + tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(32)), self.pv_acc_dtype + ) + tiled_tmem_load = tcgen05.make_tmem_copy(tmem_copy_atom, tOtO_epi) + thr_tmem_load = tiled_tmem_load.get_slice(thread_idx) + tTMEM_LOADtO = thr_tmem_load.partition_S(tOtO_epi) + tTMEM_LOADgO = thr_tmem_load.partition_D(gO_epi) + tTMEM_LOADcO = thr_tmem_load.partition_D(cO_epi) + for i in cutlass.range(cute.size(tTMEM_LOADtO, mode=[1]), unroll_full=True): + tTMEM_LOADtO_i = tTMEM_LOADtO[None, i, 0] + tTMEM_LOADgO_i = tTMEM_LOADgO[None, i, 0] + tTMEM_LOADcO_i = tTMEM_LOADcO[None, i, 0] + tTMrO = cute.make_rmem_tensor( + tTMEM_LOADcO[None, 0, i].shape, self.pv_acc_dtype + ) + cute.copy(tiled_tmem_load, tTMEM_LOADtO_i, tTMrO) + for j in cutlass.range(0, cute.size(tTMrO), 2, unroll_full=True): + tTMrO[j], tTMrO[j + 1] = cute.arch.mul_packed_f32x2( + (tTMrO[j], tTMrO[j + 1]), + (scale, scale), + ) + tSMrO = cute.make_rmem_tensor(tTMrO.shape, self.o_dtype) + o_vec = tTMrO.load() + tSMrO.store(o_vec.to(self.o_dtype)) + if cute.elem_less(tTMEM_LOADcO_i[0][0], seqlen_q): + cute.autovec_copy(tSMrO, tTMEM_LOADgO_i) + o_handle.release() + return mma_o_consumer, swap_consumer + + +def run( + q_shape: Tuple[int, int, int, int], + k_shape: Tuple[int, int, int, int], + q_dtype: Type[cutlass.Numeric], + kv_dtype: Type[cutlass.Numeric], + o_dtype: Type[cutlass.Numeric], + scale_dtype: Type[cutlass.Numeric], + scale_granularity: int, + qk_acc_dtype: Type[cutlass.Numeric], + pv_acc_dtype: Type[cutlass.Numeric], + is_persistent: bool, + is_causal: bool, + scale_q: float, + scale_k: float, + scale_v: float, + inv_scale_o: float, + scale_softmax: float, + tolerance: float, + warmup_iterations: int, + iterations: int, + skip_ref_check: bool, + use_cold_l2: bool = False, + **kwargs, +): + print(f"Running Blackwell SM100 Mixed Input FMHA Prefill D 512 test with:") + print(f" q_shape: {q_shape}") + print(f" k_shape: {k_shape}") + print(f" q_dtype: {q_dtype}") + print(f" kv_dtype: {kv_dtype}") + print(f" o_dtype: {o_dtype}") + print(f" scale_dtype: {scale_dtype}") + print(f" scale_granularity: {scale_granularity}") + print(f" qk_acc_dtype: {qk_acc_dtype}") + print(f" pv_acc_dtype: {pv_acc_dtype}") + print(f" is_persistent: {is_persistent}") + print(f" is_causal: {is_causal}") + print(f" scale_q: {scale_q}") + print(f" scale_k: {scale_k}") + print(f" scale_v: {scale_v}") + print(f" inv_scale_o: {inv_scale_o}") + print(f" scale_softmax: {scale_softmax}") + print(f" tolerance: {tolerance}") + print(f" warmup_iterations: {warmup_iterations}") + print(f" iterations: {iterations}") + print(f" skip_ref_check: {skip_ref_check}") + print(f" use_cold_l2: {use_cold_l2}") + import cutlass.torch as cutlass_torch + + # Unpack parameters + b, h_q, s_q, d = q_shape + b_, h_k, s_k, d_ = k_shape + window_size_left, window_size_right = None, None + if is_causal: + window_size_right = 0 + + if b != b_: + raise ValueError("q & k must have the same batch size") + + if d != d_: + raise ValueError("q & k must have the same head dimension") + + if d not in {512}: + raise ValueError("head dimension must be 512") + + if d % scale_granularity != 0: + raise ValueError("head dimension must be divisible by scale_granularity") + + if scale_granularity not in {128, 256, 512}: + raise ValueError("scale_granularity must be 128, 256, or 512") + + if h_q % h_k != 0: + raise ValueError("h_q must be divisible by h_k") + + if isinstance(s_q, tuple) and len(s_q) != b: + raise ValueError("variable_seqlen s_q must have the length of batch size") + if isinstance(s_k, tuple) and len(s_k) != b: + raise ValueError("variable_seqlen s_k must have the length of batch size") + + if q_dtype not in {cutlass.BFloat16}: + raise ValueError("in_dtype must be BFloat16") + + if o_dtype not in {cutlass.BFloat16}: + raise ValueError("o_dtype must be BFloat16") + + if kv_dtype not in {cutlass.Int8}: + raise ValueError("kv_dtype must be Int8") + + if qk_acc_dtype not in {cutlass.Float32}: + raise ValueError("qk_acc_dtype must be Float32") + + if pv_acc_dtype not in {cutlass.Float32}: + raise ValueError("pv_acc_dtype must be Float32") + + if not torch.cuda.is_available(): + raise RuntimeError("GPU is required to run this example!") + + torch.manual_seed(1111) + + def create_tensor(shape, dtype): + f32_torch_tensor = cutlass_torch.create_and_permute_torch_tensor( + shape, + torch.float32, + permute_order=None, + init_type=cutlass.torch.TensorInitType.RANDOM, + init_config=cutlass.torch.RandomInitConfig( + min_val=-2 if dtype.is_float or dtype.signed else 0, max_val=2 + ), + ) + + _, torch_tensor = cutlass_torch.cute_tensor_like( + f32_torch_tensor, + dtype, + is_dynamic_layout=True, + assumed_align=32, + ) + + # Create dtype cute tensor with offset (gpu) + cute_tensor = from_dlpack(torch_tensor, assumed_align=128) + cute_tensor.element_type = dtype + + return ( + f32_torch_tensor, + cute_tensor, + torch_tensor, + ) + + scale_shape = (b, h_k, s_k, d // scale_granularity) + + q_ref, q_tensor, q_torch = create_tensor(q_shape, q_dtype) + k_ref, k_tensor, k_torch = create_tensor(k_shape, kv_dtype) + v_ref, v_tensor, v_torch = create_tensor(k_shape, kv_dtype) + o_ref, o_tensor, o_torch = create_tensor(q_shape, o_dtype) + scale_k_ref, scale_k_tensor, scale_k_torch = create_tensor(scale_shape, scale_dtype) + scale_v_ref, scale_v_tensor, scale_v_torch = create_tensor(scale_shape, scale_dtype) + + mask_type = fmha_utils.MaskEnum.WINDOW_MASK_INFERENCE + if is_causal: + mask_type = fmha_utils.MaskEnum.WINDOW_MASK_INFERENCE + else: + if s_k % 128 != 0: + mask_type = fmha_utils.MaskEnum.RESIDUAL_MASK + + fmha = MixedInputFusedMultiHeadAttentionPrefillD512( + scale_granularity, + qk_acc_dtype, + pv_acc_dtype, + is_persistent, + mask_type, + ) + + # Initialize Stream + current_stream = cutlass_torch.default_stream() + + if scale_softmax == 0.0: # default to 1/sqrt(d) + scale_softmax = 1.0 / math.sqrt(d) + log2_e = math.log2( + math.exp(1.0) + ) # gpu uses exp2 for perf concerns, we need an extra factor 'log2_e' here + + scale_softmax = scale_q * scale_k * scale_softmax + scale_softmax_log2 = scale_softmax * log2_e + scale_output = scale_v * inv_scale_o + problem_size = (b, s_q, s_k, h_q, h_k, d) + compiled_fmha = cute.compile( + fmha, + q_tensor.iterator, + k_tensor.iterator, + v_tensor.iterator, + o_tensor.iterator, + scale_k_tensor.iterator, + scale_v_tensor.iterator, + problem_size, + scale_softmax_log2, + scale_output, + window_size_left if window_size_left is None else Int32(window_size_left), + window_size_right if window_size_right is None else Int32(window_size_right), + current_stream, + options=f"--opt-level 2", + ) + + def run_torch_fmha( + q, k, v, scale_k, scale_v, scale_softmax=1.0, scale_output=1.0, is_causal=False + ): + h_q = q.shape[1] + h_k = k.shape[1] + if not h_q == h_k: + repeat_factor = h_q // h_k + k = k.repeat_interleave(repeat_factor, dim=1) + v = v.repeat_interleave(repeat_factor, dim=1) + scale_k = scale_k.repeat_interleave(repeat_factor, dim=1) + scale_v = scale_v.repeat_interleave(repeat_factor, dim=1) + scale_k = ( + scale_k.unsqueeze(-1) + .repeat(1, 1, 1, 1, k.shape[3] // scale_k.shape[3]) + .reshape(k.shape) + ) + scale_v = ( + scale_v.unsqueeze(-1) + .repeat(1, 1, 1, 1, v.shape[3] // scale_v.shape[3]) + .reshape(v.shape) + ) + batch = q.shape[0] + ref_list = [] + for batch_idx in range(batch): + q_i = q[batch_idx] + k_i = k[batch_idx] + v_i = v[batch_idx] + scale_k_i = scale_k[batch_idx] + scale_v_i = scale_v[batch_idx] + s_i = torch.einsum("hqd,hkd->hqk", q_i, k_i * scale_k_i) * scale_softmax + s_q = q_i.shape[1] + s_k = k_i.shape[1] + if is_causal: + q_coords = torch.arange(0, s_q).view(-1, 1) + k_coords = torch.arange(0, s_k).view(1, -1) + _mask = k_coords > q_coords + s_k - s_q + s_i = s_i.masked_fill(_mask, -torch.inf) + p_i = s_i.softmax(dim=-1) + ref_i = torch.einsum("hqk,hkd->hqd", p_i, v_i * scale_v_i) * scale_output + ref_list.append(ref_i) + ref = torch.stack(ref_list) + return ref + + if not skip_ref_check: + # Execute kernel once for reference checking + compiled_fmha( + q_tensor.iterator, + k_tensor.iterator, + v_tensor.iterator, + o_tensor.iterator, + scale_k_tensor.iterator, + scale_v_tensor.iterator, + problem_size, + scale_softmax_log2, + scale_output, + window_size_left if window_size_left is None else Int32(window_size_left), + ( + window_size_right + if window_size_right is None + else Int32(window_size_right) + ), + current_stream, + ) + print("Verifying results...") + o_ref = run_torch_fmha( + q_ref, + k_ref, + v_ref, + scale_k_ref, + scale_v_ref, + scale_softmax, + scale_output, + is_causal, + ) + + # convert o back to f32 for comparison + o_fp32, o_fp32_torch = cutlass_torch.cute_tensor_like( + torch.empty(*o_torch.shape, dtype=torch.float32), + Float32, + is_dynamic_layout=True, + assumed_align=16, + ) + cute.testing.convert(o_tensor, o_fp32) + o_result = o_fp32_torch.cpu() + torch.testing.assert_close(o_ref, o_result, atol=tolerance, rtol=1e-05) + + print("Results verified successfully!") + + +if __name__ == "__main__": + + def parse_comma_separated_ints(s: str): + try: + return tuple(int(x.strip()) for x in s.split(",")) + except ValueError: + raise argparse.ArgumentTypeError( + "Invalid format. Expected comma-separated integers." + ) + + parser = argparse.ArgumentParser(description="Example of FMHA on Blackwell.") + + parser.add_argument( + "--q_dtype", + type=cutlass.dtype, + default=cutlass.BFloat16, + help="Input data type", + ) + + parser.add_argument( + "--kv_dtype", + type=cutlass.dtype, + default=cutlass.Int8, + help="Input data type", + ) + + parser.add_argument( + "--o_dtype", + type=cutlass.dtype, + default=cutlass.BFloat16, + help="Output data type", + ) + + parser.add_argument( + "--scale_dtype", + type=cutlass.dtype, + default=cutlass.BFloat16, + help="Scale data type", + ) + + parser.add_argument( + "--scale_granularity", + type=int, + default=512, + help="Scale granularity", + ) + + parser.add_argument( + "--qk_acc_dtype", + type=cutlass.dtype, + default=Float32, + help="QK accumulator data type", + ) + + parser.add_argument( + "--pv_acc_dtype", + type=cutlass.dtype, + default=Float32, + help="PV accumulator data type", + ) + + parser.add_argument( + "--is_persistent", + action="store_true", + help="Is persistent", + ) + + parser.add_argument( + "--is_causal", + action="store_true", + help="Whether to use casual mask", + ) + + parser.add_argument( + "--q_shape", + type=parse_comma_separated_ints, + default=(1, 8, 256, 512), + help="Shape of Q (B, H, S_q, D)", + ) + + parser.add_argument( + "--k_shape", + type=parse_comma_separated_ints, + default=(1, 8, 256, 512), + help="Shape of K (B, H_k, S_k, D)", + ) + + parser.add_argument( + "--scale_q", + type=float, + default=1.0, + help="Scaling factors to dequantize Q", + ) + + parser.add_argument( + "--scale_k", + type=float, + default=1.0, + help="Scaling factors to dequantize K", + ) + + parser.add_argument( + "--scale_v", + type=float, + default=1.0, + help="Scaling factors to dequantize V", + ) + + parser.add_argument( + "--inv_scale_o", + type=float, + default=1.0, + help="Scaling factor to quantize O", + ) + + parser.add_argument( + "--scale_softmax", + type=float, + default=0.0, + help="Scaling factor to scale S (i.e. Q*K); if zero, defaults to 1/sqrt(D)", + ) + + parser.add_argument( + "--tolerance", type=float, default=1e-01, help="Tolerance for validation" + ) + + parser.add_argument( + "--warmup_iterations", + type=int, + default=0, + help="Number of iterations for warmup", + ) + + parser.add_argument( + "--iterations", + type=int, + default=1, + help="Number of iterations after warmup", + ) + + parser.add_argument( + "--skip_ref_check", + action="store_true", + help="Skip reference check", + ) + + parser.add_argument( + "--use_cold_l2", + action="store_true", + default=False, + help="Use circular buffer tensor sets to ensure L2 cold cache", + ) + + args = parser.parse_args() + + if len(args.q_shape) != 4: + parser.error("--q_shape must contain exactly 4 values") + + if len(args.k_shape) != 4: + parser.error("--k_shape must contain exactly 4 values") + + if not torch.cuda.is_available(): + raise RuntimeError("GPU is required to run this example!") + + torch.manual_seed(1111) + + run( + args.q_shape, + args.k_shape, + args.q_dtype, + args.kv_dtype, + args.o_dtype, + args.scale_dtype, + args.scale_granularity, + args.qk_acc_dtype, + args.pv_acc_dtype, + args.is_persistent, + args.is_causal, + args.scale_q, + args.scale_k, + args.scale_v, + args.inv_scale_o, + args.scale_softmax, + args.tolerance, + args.warmup_iterations, + args.iterations, + args.skip_ref_check, + args.use_cold_l2, + ) + + print("PASS") diff --git a/examples/python/CuTeDSL/blackwell/mixed_input_fmha/prefill_helpers.py b/examples/python/CuTeDSL/blackwell/mixed_input_fmha/prefill_helpers.py new file mode 100644 index 00000000..2a833cce --- /dev/null +++ b/examples/python/CuTeDSL/blackwell/mixed_input_fmha/prefill_helpers.py @@ -0,0 +1,400 @@ +# 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. + +from typing import Tuple, Optional + +import cutlass +import cutlass.cute as cute +import cutlass.cute.nvgpu.tcgen05 as tcgen05 +import cutlass.pipeline as pipeline + + +@cute.jit +def load_qk( + iterations: int, + kv_step: cutlass.Int32, + k_args: Tuple, + scale_k_args: Optional[Tuple] = None, + q_args: Optional[Tuple] = None, +) -> Tuple[pipeline.PipelineProducer, pipeline.PipelineProducer]: + if cutlass.const_expr(q_args is not None): + tQgQ, tQsQ, tma_atom_q, load_q_producer = q_args + else: + tQgQ, tQsQ, tma_atom_q, load_q_producer = None, None, None, None + tKgK, tKsK, tma_atom_k, load_k_producer = k_args + tKgScaleK, tKsScaleK, tma_atom_scale_k, load_scale_k_producer = scale_k_args + + scale_k_handle = load_scale_k_producer.acquire_and_advance() + cute.copy( + tma_atom_scale_k, + tKgScaleK[None, kv_step], + tKsScaleK[None, scale_k_handle.index], + tma_bar_ptr=scale_k_handle.barrier, + ) + for iter in cutlass.range(iterations, unroll=1): + if cutlass.const_expr(q_args is not None): + q_handle = load_q_producer.acquire_and_advance() + cute.copy( + tma_atom_q, + tQgQ[None, iter], + tQsQ[None, q_handle.index], + tma_bar_ptr=q_handle.barrier, + ) + k_handle = load_k_producer.acquire_and_advance() + cute.copy( + tma_atom_k, + tKgK[None, kv_step, iter], + tKsK[None, k_handle.index], + tma_bar_ptr=k_handle.barrier, + ) + if cutlass.const_expr(q_args is not None): + return load_k_producer, load_scale_k_producer, load_q_producer + else: + return load_k_producer, load_scale_k_producer + + +@cute.jit +def load_v( + iterations: int, + kv_step: cutlass.Int32, + v_args: Tuple, + scale_v_args: Tuple, +) -> pipeline.PipelineProducer: + tVgV, tVsV, tma_atom_v, load_v_producer = v_args + tScaleVgV, tScaleVsV, tma_atom_scale_v, load_scale_v_producer = scale_v_args + scale_v_handle = load_scale_v_producer.acquire_and_advance() + cute.copy( + tma_atom_scale_v, + tScaleVgV[None, kv_step], + tScaleVsV[None, scale_v_handle.index], + tma_bar_ptr=scale_v_handle.barrier, + ) + for iter in cutlass.range(iterations, unroll=1): + v_handle = load_v_producer.acquire_and_advance() + cute.copy( + tma_atom_v, + tVgV[None, iter, kv_step], + tVsV[None, v_handle.index], + tma_bar_ptr=v_handle.barrier, + ) + return load_v_producer, load_scale_v_producer + + +@cute.jit +def get_scale_smem_layout( + scale_granularity: int, + d_r: int, + mma_tiler: cute.Tile, + major_mode: tcgen05.OperandMajorMode, +) -> Tuple[cute.Layout, cute.Tile]: + size_mn = mma_tiler[1] // 2 # 2cta by default + if cutlass.const_expr(major_mode == tcgen05.OperandMajorMode.MN): # v + scale_tiler = (mma_tiler[2] * d_r,) + tma_view_layout = cute.make_layout( + (mma_tiler[2] * d_r), + ) + assert scale_granularity % mma_tiler[1] == 0, ( + "scale_granularity must be divisible by mma_tiler[1]" + ) + rest_l = scale_granularity // mma_tiler[1] + s2r_view_layout = cute.make_layout( + (size_mn, mma_tiler[2], (rest_l, d_r)), + stride=(0, d_r, (0, 1)), + ) + else: # k + scale_tiler = (mma_tiler[1] * d_r,) + tma_view_layout = cute.make_layout((size_mn * d_r)) + assert scale_granularity % mma_tiler[2] == 0, ( + "scale_granularity must be divisible by mma_tiler[2]" + ) + rest_l = scale_granularity // mma_tiler[2] + s2r_view_layout = cute.make_layout( + (size_mn, mma_tiler[2], (rest_l, d_r)), + stride=(d_r, 0, (0, 1)), + ) + # Apply a trivial swizzle to make it a composed layout, which could be used to construct TMA atom + tma_view_smem_layout = cute.make_composed_layout( + cute.make_swizzle(0, 4, 3), 0, tma_view_layout + ) + return tma_view_smem_layout, scale_tiler, s2r_view_layout + + +@cute.jit +def mma_qk( + iterations: int, + qk_tiled_mma: cute.TiledMma, + tensor_args: Tuple, + pipeline_args: Tuple, +): + tStS, tSrQ, tSrK_trans = tensor_args + mma_s_producer, load_q_consumer, load_q_releaser, dequant_kv_consumer = ( + pipeline_args + ) + cta_rank_in_cluster = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) + is_leader_cta = cta_rank_in_cluster % 2 == 0 + if is_leader_cta: + s_handle = mma_s_producer.acquire_and_advance() + tStS_slice = tStS[None, None, None, s_handle.index] + qk_tiled_mma.set(tcgen05.Field.ACCUMULATE, False) + for iter in cutlass.range(iterations, unroll=1): + if cutlass.const_expr(load_q_consumer is not None): + load_q_consumer.wait_and_advance() + tSrQ_slice = tSrQ[None, None, None, iter] + k_trans_handle = dequant_kv_consumer.wait_and_advance() + tSrK_trans_slice = tSrK_trans[None, None, None, k_trans_handle.index] + num_kphases = cute.size(tSrQ_slice, mode=[2]) + for kphase_idx in cutlass.range(num_kphases, unroll_full=True): + kphase_coord = (None, None, kphase_idx) + cute.gemm( + qk_tiled_mma, + tStS_slice, + tSrQ_slice[kphase_coord], + tSrK_trans_slice[kphase_coord], + tStS_slice, + ) + qk_tiled_mma.set(tcgen05.Field.ACCUMULATE, True) + k_trans_handle.release() + if cutlass.const_expr(load_q_releaser is not None): + load_q_releaser.release() + load_q_releaser.advance() + s_handle.commit() + return mma_s_producer, load_q_consumer, dequant_kv_consumer + + +@cute.jit +def dequant_k( + iterations: int, + transform_warp_ids: Tuple, + dtype_args: Tuple, + tensor_args: Tuple, + pipeline_args: Tuple, +): + (k_dtype, q_dtype) = dtype_args + (sOrig, sScale, sTrans) = tensor_args + (load_kv_consumer, load_scale_consumer, dequant_kv_producer) = pipeline_args + tidx, _, _ = cute.arch.thread_idx() + THREADS_PER_WARP = 32 + thread_idx = tidx % (THREADS_PER_WARP * len(transform_warp_ids)) + r2s_copy_atom = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), k_dtype, num_bits_per_copy=32 + ) + # Construct tiled_copy satisfying 16 contiguous elts per copy atom + r2s_tiled_copy = cute.make_cotiled_copy( + r2s_copy_atom, + cute.make_layout((256, 16), stride=(16, 1)), + sTrans[(None, None, None, 0)].layout, + ) + thr_r2s_tiled_copy = r2s_tiled_copy.get_slice(thread_idx) + tOsOrig = thr_r2s_tiled_copy.partition_S(sOrig) + tTsTrans = thr_r2s_tiled_copy.partition_D(sTrans) + tOrOrig = cute.make_rmem_tensor_like( + cute.append( + tOsOrig[None, None, None, None, 0].layout, + cute.make_layout( + 2, stride=cute.cosize(tOsOrig[None, None, None, None, 0].layout) + ), + ), + k_dtype, + ) + tTrTrans = cute.make_rmem_tensor_like( + cute.append( + tTsTrans[None, None, None, None, 0].layout, + cute.make_layout( + 2, stride=cute.cosize(tTsTrans[None, None, None, None, 0].layout) + ), + ), + q_dtype, + ) + tSsScale = thr_r2s_tiled_copy.partition_S(sScale) + tSrScale = cute.make_rmem_tensor_like(tSsScale[None, None, None, None, None, 0]) + scale_handle = load_scale_consumer.wait_and_advance() + cute.autovec_copy( + tSsScale[None, None, None, None, None, scale_handle.index], tSrScale + ) + cute.arch.fence_view_async_shared() + scale_handle.release() + # prefetch iter = 0 + kv_handle = load_kv_consumer.wait_and_advance() + cute.autovec_copy( + tOsOrig[None, None, None, None, kv_handle.index], + tOrOrig[None, None, None, None, 0], + ) + transformed_tensor = tOrOrig[None, None, None, None, 0].load().to(q_dtype) + scale = cute.TensorSSA( + tSrScale[None, None, None, None, 0].load(), + transformed_tensor.shape, + q_dtype, + ) + transformed_tensor = transformed_tensor * scale + tTrTrans[None, None, None, None, 0].store(transformed_tensor) + cute.arch.fence_view_async_shared() + kv_handle.release() + for iter in cutlass.range(1, iterations, unroll_full=True): + kv_trans_handle = dequant_kv_producer.acquire_and_advance() + cute.autovec_copy( + tTrTrans[None, None, None, None, (iter - 1) % 2], + tTsTrans[None, None, None, None, kv_trans_handle.index], + ) + cute.arch.fence_view_async_shared() + kv_trans_handle.commit() + kv_handle = load_kv_consumer.wait_and_advance() + cute.autovec_copy( + tOsOrig[None, None, None, None, kv_handle.index], + tOrOrig[None, None, None, None, iter % 2], + ) + transformed_tensor = ( + tOrOrig[None, None, None, None, iter % 2].load().to(q_dtype) + ) + scale = cute.TensorSSA( + tSrScale[None, None, None, None, iter].load(), + transformed_tensor.shape, + q_dtype, + ) + transformed_tensor = transformed_tensor * scale + tTrTrans[None, None, None, None, iter % 2].store(transformed_tensor) + cute.arch.fence_view_async_shared() + kv_handle.release() + kv_trans_handle = dequant_kv_producer.acquire_and_advance() + cute.autovec_copy( + tTrTrans[None, None, None, None, (iterations - 1) % 2], + tTsTrans[None, None, None, None, kv_trans_handle.index], + ) + cute.arch.fence_view_async_shared() + kv_trans_handle.commit() + return load_kv_consumer, load_scale_consumer, dequant_kv_producer + + +@cute.jit +def dequant_v( + iterations: int, + transform_warp_ids: Tuple, + dtype_args: Tuple, + tensor_args: Tuple, + pipeline_args: Tuple, +): + (v_dtype, q_dtype) = dtype_args + (sOrig, sScale, sTrans) = tensor_args + (load_kv_consumer, load_scale_consumer, dequant_kv_producer) = pipeline_args + tidx, _, _ = cute.arch.thread_idx() + THREADS_PER_WARP = 32 + thread_idx = tidx % (THREADS_PER_WARP * len(transform_warp_ids)) + r2s_copy_atom = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), v_dtype, num_bits_per_copy=32 + ) + # Construct tiled_copy satisfying 16 contiguous elts per copy atom + r2s_tiled_copy = cute.make_cotiled_copy( + r2s_copy_atom, + cute.make_layout((256, 16), stride=(16, 1)), + sTrans[(None, None, None, 0)].layout, + ) + thr_r2s_tiled_copy = r2s_tiled_copy.get_slice(thread_idx) + tOsOrig = thr_r2s_tiled_copy.partition_S(sOrig) + tTsTrans = thr_r2s_tiled_copy.partition_D(sTrans) + # double buffer for better perf + tOrOrig = cute.make_rmem_tensor_like( + cute.append( + tOsOrig[None, None, None, None, 0].layout, + cute.make_layout( + 2, stride=cute.cosize(tOsOrig[None, None, None, None, 0].layout) + ), + ), + v_dtype, + ) + tTrTrans = cute.make_rmem_tensor_like( + cute.append( + tTsTrans[None, None, None, None, 0].layout, + cute.make_layout( + 2, stride=cute.cosize(tTsTrans[None, None, None, None, 0].layout) + ), + ), + q_dtype, + ) + tSsScale = thr_r2s_tiled_copy.partition_S(sScale) + tSrScale = cute.make_rmem_tensor_like(tSsScale[None, None, None, None, None, 0]) + scale_v_handle = load_scale_consumer.wait_and_advance() + cute.autovec_copy( + tSsScale[None, None, None, None, None, scale_v_handle.index], + tSrScale, + ) + cute.arch.fence_view_async_shared() + scale_v_handle.release() + # prefetch iter = 0 + kv_handle = load_kv_consumer.wait_and_advance() + cute.autovec_copy( + tOsOrig[None, None, None, None, kv_handle.index], + tOrOrig[None, None, None, None, 0], + ) + transformed_tensor = tOrOrig[None, None, None, None, 0].load().to(q_dtype) + scale = cute.TensorSSA( + tSrScale[None, None, None, None, 0].load(), + transformed_tensor.shape, + q_dtype, + ) + transformed_tensor = transformed_tensor * scale + tTrTrans[None, None, None, None, 0].store(transformed_tensor) + cute.arch.fence_view_async_shared() + kv_handle.release() + for iter in cutlass.range(1, iterations, unroll_full=True): + kv_trans_handle = dequant_kv_producer.acquire_and_advance() + cute.autovec_copy( + tTrTrans[None, None, None, None, (iter - 1) % 2], + tTsTrans[None, None, None, None, kv_trans_handle.index], + ) + cute.arch.fence_view_async_shared() + kv_trans_handle.commit() + kv_handle = load_kv_consumer.wait_and_advance() + cute.autovec_copy( + tOsOrig[None, None, None, None, kv_handle.index], + tOrOrig[None, None, None, None, iter % 2], + ) + transformed_tensor = ( + tOrOrig[None, None, None, None, iter % 2].load().to(q_dtype) + ) + scale = cute.TensorSSA( + tSrScale[ + None, + None, + None, + None, + iter, + ].load(), + transformed_tensor.shape, + q_dtype, + ) + transformed_tensor = transformed_tensor * scale + tTrTrans[None, None, None, None, iter % 2].store(transformed_tensor) + cute.arch.fence_view_async_shared() + kv_handle.release() + kv_trans_handle = dequant_kv_producer.acquire_and_advance() + cute.autovec_copy( + tTrTrans[None, None, None, None, (iterations - 1) % 2], + tTsTrans[None, None, None, None, kv_trans_handle.index], + ) + cute.arch.fence_view_async_shared() + kv_trans_handle.commit() + return load_kv_consumer, load_scale_consumer, dequant_kv_producer diff --git a/examples/python/CuTeDSL/blackwell/mixed_input_gemm/grouped_mixed_input_gemm.py b/examples/python/CuTeDSL/blackwell/mixed_input_gemm/grouped_mixed_input_gemm.py new file mode 100644 index 00000000..a81d1e39 --- /dev/null +++ b/examples/python/CuTeDSL/blackwell/mixed_input_gemm/grouped_mixed_input_gemm.py @@ -0,0 +1,2526 @@ +# 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. + +import argparse +from math import log2, ceil +from typing import Optional +import os +import sys + +import cuda.bindings.driver as cuda + +import cutlass +import cutlass.cute as cute +import cutlass.pipeline as pipeline +from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait +import cutlass.utils as utils +import cutlass.utils.blackwell_helpers as sm100_utils +import cutlass.utils.mixed_input_helpers as mixed_input_utils +from cutlass.utils.mixed_input_helpers import TransformMode +import cutlass.cute.testing as testing +from cutlass.cute.nvgpu import cpasync, tcgen05 + +if __name__ == "__main__": + current_dir = os.path.dirname(os.path.abspath(__file__)) + sys.path.insert(0, os.path.join(current_dir, "../..")) + +from blackwell.mixed_input_gemm.mixed_input_host_utils import ( + create_tensors_for_contiguous_grouped_mixed_input_gemm as create_tensors, + run_contiguous_grouped_ref_and_compare as run_ref_and_compare, +) +from cutlass.cute.runtime import from_dlpack + +""" +A mixed-input grouped GEMM example for the NVIDIA Blackwell SM100 architecture using CUTE DSL. + +This example demonstrates an implementation of mixed-input grouped GEMM using a TMA plus Blackwell +SM100 TensorCore warp-specialized persistent kernel. It can be viewed as an extension of the batched +mixed-input GEMM example to support a specific grouped GEMM pattern: grouped GEMM with contiguous offsets. + +Specifically, the input A tensor is still in the shape of (M, K, L), and L is the number of groups. The +input B tensor is in the shape of (N, K) and the result C tensor is in the shape of (M, N). Tensor B +and tensor C are not divided into groups explicitly and there is an extra input tensor cumsum defining +the mapping between the N mode to groups. The cumsum tensor is in the shape of (N+1) and cumsum[i] +defines the accumulated size along N mode for groups up to i (not including i): + + ``` + Group 0 Group 1 Group 2 ..... Group L-1 + -+--------+--------+--------+.....+----------------+ + | | | | | + |<- N0 ->|<- N1 ->|<- N2 ->|.....|<-- NL-1 -->| + | | | | | + -+--------+--------+--------+.....+-------------------+ +cumsum: | 0 | N0 | N0+N1 |.....| sum(N0,N1,...NL-2) | sum(N0,N1,...NL-1) + ``` + +The computation flow is the same as the batched mixed-input GEMM example. A is the narrow-precision tensor +and B holds data with a wider precision. MMA will work in the wide precision of tensor B and tensor A +will be transformed to the wide precision of tensor B following 1 of the 2 possible modes as follows: + +1. convert-only mode: + C = type_convert(A) x B + +In convert-only mode, tensor A is directly converted to the wide precision of tensor B. + +2. convert-scale mode: + C = (type_convert(A) * scale) x B + +In convert-scale mode, tensor A is first converted to the wide precision of tensor B and then scaled by the scale tensor. +The scale tensor is in the same precision as tensor B. +The mode is determined by tensor A's data type as follows: +- if tensor A is in int8 or uint8, convert-only mode is used. +- if tensor A is in int4, convert-scale mode is used. + +The output tensor C could have the same precision as tensor B or fp32. + +To run this example: + +.. code-block:: bash + + python examples/blackwell/grouped_mixed_input_gemm.py \ + --a_dtype Int8 --b_dtype BFloat16 \ + --scale_granularity_m 0 --scale_granularity_k 0 \ + --c_dtype BFloat16 --acc_dtype Float32 \ + --mma_tiler_mnk 128,128,64 --cluster_shape_mn 1,1 \ + --mnkl 256,512,8192,1 + +Input A and B have int8 and bf16 data types, respectively. The Blackwell tcgen05 MMA tile shape +is specified as (128,128,64) and the cluster shape is (1,1). The MMA accumulator and output data type +are set as fp32 and bf16, respectively. As tensor A is int8, convert-only mode is used. +scale_granularity_m and scale_granularity_k are set as 0 for convert-only mode. + +Here is an example of running convert-scale mode: + +.. code-block:: bash + + python examples/blackwell/mixed_input_gemm/grouped_mixed_input_gemm.py \ + --a_dtype Int4 --b_dtype BFloat16 \ + --scale_granularity_m 1 --scale_granularity_k 256 \ + --c_dtype BFloat16 --acc_dtype Float32 \ + --mma_tiler_mnk 256,128,128 --cluster_shape_mn 2,1 \ + --use_2cta_instrs --mnkl 1024,8192,6144,16 \ + +Input A and B have int4 and bf16 data types, respectively. The scale granularity is set as (1,256), +which means each element along the m mode of tensor A has its own scale element and 256 contiguous elements +along the k mode share the same scale element. There is no scale reuse along the L mode. If the GEMM shape is +(M, N, K, L), then the scale tensor shape is (M // scale_granularity_m, K // scale_granularity_k, L), +which is (1024, 6144/256, 16) in this example. +The Blackwell tcgen05 MMA tile shape is specified as (256,128,128) and tcgen05 2CTA feature is enabled. +The cluster shape is (2,1). The MMA accumulator and output data type are set as fp32 and bf16, respectively. +As tensor A is int4, the convert-scale mode is used. + +To collect performance with NCU profiler: + +.. code-block:: bash + + ncu python examples/blackwell/mixed_input_gemm/grouped_mixed_input_gemm.py \ + --a_dtype Int8 --b_dtype BFloat16 \ + --scale_granularity_m 0 --scale_granularity_k 0 \ + --c_dtype BFloat16 --acc_dtype Float32 \ + --mma_tiler_mnk 128,128,64 --cluster_shape_mn 1,1 \ + --mnkl 256,512,8192,1 \ + --warmup_iterations 1 --iterations 10 --skip_ref_check + +Besides the requirements from the batched mixed-input GEMM example, there are some constraints for this example: +* --use_tma_store option is removed as no alignment assumption is made for each group. +""" + + +class GroupedMixedInputGemmKernel: + """ + Mixed-input grouped GEMM kernel for NVIDIA Blackwell SM100 architecture. + + This kernel supports GEMM operations where input tensors A and B have different + data types, with tensor A being transformed to the precision of tensor B before + matrix multiplication. + Tensor A is in shape of [M, K, L] with L being the number of groups. Tensor B is in shape of [N, K] and a group search algorithm + is applied along the N mode to find the group index for each CTA tile. A cumsum tensor provides the offsets of each group along the N mode. + + :param scale_granularity_m: Number of elements sharing the same scale factor along the M mode + :type scale_granularity_m: int + :param scale_granularity_k: Number of elements sharing the same scale factor along the K mode + :type scale_granularity_k: int + :param acc_dtype: Data type for accumulation during computation + :type acc_dtype: type[cutlass.Numeric] + :param use_2cta_instrs: Whether to use CTA group 2 for advanced thread cooperation + :type use_2cta_instrs: bool + :param mma_tiler_mnk: Shape of the Matrix Multiply-Accumulate (MMA) tile (M, N, K) + :type mma_tiler_mnk: tuple[int, int, int] + :param cluster_shape_mn: Cluster dimensions (M,N) for parallel processing + :type cluster_shape_mn: tuple[int, int] + :param group_count: The total number of groups + :type group_count: int + :param shuffle_a: Whether to use shuffle intrinsic for int4-to-bf16 conversion + :type shuffle_a: bool + """ + + def __init__( + self, + scale_granularity_m: int, + scale_granularity_k: int, + acc_dtype: type[cutlass.Numeric], + use_2cta_instrs: bool, + mma_tiler_mnk: tuple[int, int, int], + cluster_shape_mn: tuple[int, int], + group_count: int, + shuffle_a: bool, + ): + """ + Initializes the mixed-input GEMM kernel with a specified configuration. + """ + # Scale granularity defines how many elements share the same scale factor + # along the M and K modes. + self.scale_granularity_m = scale_granularity_m + self.scale_granularity_k = scale_granularity_k + # Set transform mode + if cutlass.const_expr( + self.scale_granularity_m == 0 and self.scale_granularity_k == 0 + ): + self.scale_mode = TransformMode.ConvertOnly + else: + self.scale_mode = TransformMode.ConvertScale + self.group_count = group_count + self.acc_dtype = acc_dtype + self.use_2cta_instrs = use_2cta_instrs + self.cluster_shape_mn = cluster_shape_mn + self.mma_tiler = mma_tiler_mnk + self.shuffle_a = shuffle_a + self.cta_group = ( + tcgen05.CtaGroup.TWO if self.use_2cta_instrs else tcgen05.CtaGroup.ONE + ) + # Set specialized warp ids + self.epilog_warp_id = ( + 0, + 1, + 2, + 3, + ) + self.mma_warp_id = 4 + self.tma_warp_id = 5 + self.scale_tma_warp_id = 6 + # Schedule warp to do the group search + self.schedule_warp_id = 7 + self.transform_warp_id = ( + 8, + 9, + 10, + 11, + ) + # Define expected register count for different warps + self.num_regs_epilogue_warps = 192 + self.num_regs_mma_warp = 96 + self.num_regs_tma_warps = 80 + self.num_regs_transform_warps = 208 + self.num_regs_schedule_warp = 64 + self.threads_per_cta = 32 * ( + max( + ( + self.mma_warp_id, + self.tma_warp_id, + self.scale_tma_warp_id, + *self.epilog_warp_id, + *self.transform_warp_id, + ) + ) + + 1 + ) + + # Set barrier id for cta sync, epilogue sync, tmem ptr sync, and transform sync + self.epilog_sync_barrier = pipeline.NamedBarrier( + 1, 32 * len(self.epilog_warp_id) + ) + self.tmem_ptr_sync_barrier = pipeline.NamedBarrier(2, self.threads_per_cta) + self.transform_sync_barrier = pipeline.NamedBarrier( + 3, 32 * len(self.transform_warp_id) + ) + self.cta_sync_barrier = pipeline.NamedBarrier(4, self.threads_per_cta) + self.sched_sync_barrier = pipeline.NamedBarrier(5, 32) + + self.smem_buffer_align_bytes = 1024 + + def _setup_attributes(self): + """Set up configurations that are dependent on GEMM inputs + + This method configures various attributes based on the input tensor properties + (data types, leading dimensions) and kernel settings: + - Deduce where the transformed A tensor is stored + - Configuring tiled MMA + - Computing MMA/cluster/tile shapes + - Computing cluster layout + - Computing multicast CTAs for A/B + - Computing epilogue sub-tile + - Setting up A/scale/B/C stage counts in shared memory + - Setting up transformed A stage count in shared memory or tensor memory + - Computing A/transformed A/scale/B/C memory layout + - Computing tensor memory allocation columns + """ + # Deduce where the transformed A tensor is stored, shared memory(SMEM) or tensor memory(TMEM) + self.transform_a_source = mixed_input_utils.get_transform_a_source( + self.a_major_mode + ) + tiled_mma = sm100_utils.make_trivial_tiled_mma( + self.mma_dtype, + self.a_major_mode, + self.b_major_mode, + self.acc_dtype, + self.cta_group, + self.mma_tiler[:2], + self.transform_a_source, + ) + self.cta_tile_shape_mnk = ( + self.mma_tiler[0] // cute.size(tiled_mma.thr_id.shape), + self.mma_tiler[1], + self.mma_tiler[2], + ) + self.cluster_tile_shape_mnk = ( + self.cluster_shape_mn[0] * self.cta_tile_shape_mnk[0], + self.cluster_shape_mn[1] * self.cta_tile_shape_mnk[1], + self.cta_tile_shape_mnk[2], + ) + self.cluster_layout_vmnk = cute.tiled_divide( + cute.make_layout((*self.cluster_shape_mn, 1)), + (tiled_mma.thr_id.shape,), + ) + self.num_mcast_ctas_a = cute.size(self.cluster_layout_vmnk.shape[2]) + self.num_mcast_ctas_b = cute.size(self.cluster_layout_vmnk.shape[1]) + self.is_a_mcast = self.num_mcast_ctas_a > 1 + self.is_b_mcast = self.num_mcast_ctas_b > 1 + + self.epi_tile = sm100_utils.compute_epilogue_tile_shape( + self.cta_tile_shape_mnk, + self.use_2cta_instrs, + self.c_layout, + self.c_dtype, + ) + + # Compute tensor memory(TMEM) columns and stages for each pipeline + ( + self.num_load2trans_stage, + self.num_scale_load2trans_stage, + self.num_trans2mma_stage, + self.num_acc_stage, + self.num_c_stage, + self.num_tile_info_stage, + self.num_acc_tmem_cols, + self.num_a_tmem_cols, + ) = self._compute_stages_and_tmem_cols( + tiled_mma, + self.mma_tiler, + self.cta_tile_shape_mnk, + self.epi_tile, + self.a_dtype, + self.b_dtype, + self.c_dtype, + self.c_layout, + self.transform_a_source, + self.scale_granularity_m, + self.scale_granularity_k, + self.smem_buffer_align_bytes, + self.scale_mode, + ) + + # Align TMEM columns for allocation + # TMEM allocation requires power-of-2 column alignment + # and must meet minimum allocation requirements + self.num_tmem_alloc_cols = cute.round_up( + self.num_acc_tmem_cols + self.num_a_tmem_cols, + cute.arch.get_min_tmem_alloc_cols("sm_100"), + ) + self.num_tmem_alloc_cols = 2 ** (ceil(log2(self.num_tmem_alloc_cols))) + # Get smem layout for C tensor + self.c_smem_layout_staged = sm100_utils.make_smem_layout_epi( + self.c_dtype, + self.c_layout, + self.epi_tile, + self.num_c_stage, + ) + # Get smem layout for A, transformed A, and B + ( + self.smem_layout_a, + self.smem_layout_a_transform, + self.smem_layout_b, + ) = mixed_input_utils.compute_smem_layout( + tiled_mma, + self.mma_tiler, + self.a_dtype, + self.b_dtype, + self.num_load2trans_stage, + self.num_trans2mma_stage, + ) + # Get smem layout for scale tensor + self.smem_layout_scale_per_stage = None + self.smem_layout_scale = None + if cutlass.const_expr(self.scale_mode == TransformMode.ConvertScale): + # Get scale tile shape and smem layout for scale tensor + ( + self.scale_tile_shape, + self.smem_layout_scale_per_stage, + self.smem_layout_scale, + ) = mixed_input_utils.get_smem_layout_scale( + self.mma_tiler, + self.use_2cta_instrs, + self.scale_granularity_m, + self.scale_granularity_k, + self.scale_major_mode, + self.a_scale_dtype, + self.num_scale_load2trans_stage, + ) + + def _validate_inputs( + self, + a: cute.Tensor, + a_scale: Optional[cute.Tensor], + b: cute.Tensor, + c: cute.Tensor, + ) -> None: + """ + Validates input tensors and their properties. + + :param a: Input tensor A. + :type a: cute.Tensor + :param a_scale: Scale tensor for tensor A (None for ConvertOnly mode). + :type a_scale: Optional[cute.Tensor] + :param b: Input tensor B. + :type b: cute.Tensor + :param c: Output tensor C. + :type c: cute.Tensor + :raises ValueError: If inputs don't meet kernel requirements. + """ + # Validate scale tensor major mode + if cutlass.const_expr( + self.scale_mode == TransformMode.ConvertScale + and utils.LayoutEnum.from_tensor(a_scale).mma_major_mode() + != tcgen05.OperandMajorMode.MN + ): + raise ValueError("scale_major_mode must be M-major") + + @cute.jit + def __call__( + self, + a: cute.Tensor, + a_scale: Optional[cute.Tensor], # None for ConvertOnly mode + b: cute.Tensor, + cumsum: cute.Tensor, + c: cute.Tensor, + max_active_clusters: cutlass.Constexpr, + stream: cuda.CUstream, + ): + """ + Executes the Mixed Input Grouped GEMM operation. + + This method sets up the kernel parameters, computes the grid size, + defines the shared storage, and launches the kernel. + + The execution steps are as follows: + - Setup static attributes before smem/grid/tma computation. + - Setup TMA load/store atoms and tensors. + - Compute grid size with regard to hardware constraints. + - Define shared storage for kernel. + - Launch the kernel synchronously. + + :param a: Input tensor A. + :type a: cute.Tensor + :param a_scale: Scale tensor for tensor A (None for ConvertOnly mode). + :type a_scale: Optional[cute.Tensor] + :param b: Input tensor B. + :type b: cute.Tensor + :param cumsum: tensor containing the cumulative size of each group along the search mode(aka, N mode in this example). + :type cumsum: cute.Tensor + :param c: Output tensor C. + :type c: cute.Tensor + :param max_active_clusters: Maximum number of active clusters to launch. + :type max_active_clusters: cutlass.Constexpr + :param stream: CUDA stream to launch the kernel on. + :type stream: cuda.CUstream + """ + self.a_dtype: type[cutlass.Numeric] = a.element_type + self.a_scale_dtype: type[cutlass.Numeric] = ( + a_scale.element_type + if self.scale_mode is TransformMode.ConvertScale + else None + ) + self.b_dtype: type[cutlass.Numeric] = b.element_type + self.c_dtype: type[cutlass.Numeric] = c.element_type + self.mma_dtype = self.b_dtype + + self.a_major_mode = utils.LayoutEnum.from_tensor(a).mma_major_mode() + self.scale_major_mode = ( + utils.LayoutEnum.from_tensor(a_scale).mma_major_mode() + if self.scale_mode is TransformMode.ConvertScale + else None + ) + self.b_major_mode = utils.LayoutEnum.from_tensor(b).mma_major_mode() + self.c_layout = utils.LayoutEnum.from_tensor(c) + if cutlass.const_expr(self.scale_mode == TransformMode.ConvertScale): + # Get gmem layout for scale tensor + self.gmem_layout_scale = mixed_input_utils.get_gmem_layout_scale( + a.shape, + self.scale_granularity_m, + self.scale_granularity_k, + self.scale_major_mode, + ) + + # Validate inputs + self._validate_inputs(a, a_scale, b, c) + + # Setup attributes that dependent on gemm inputs + self._setup_attributes() + + tiled_mma = sm100_utils.make_trivial_tiled_mma( + self.mma_dtype, + self.a_major_mode, + self.b_major_mode, + self.acc_dtype, + self.cta_group, + self.mma_tiler[:2], + self.transform_a_source, + ) + # Set up gmem copy atoms for A, scale, and B + a_op = mixed_input_utils.get_tma_atom_kind( + self.is_a_mcast, self.use_2cta_instrs, False + ) + b_op = mixed_input_utils.get_tma_atom_kind( + self.is_b_mcast, self.use_2cta_instrs, True + ) + a_scale_op = a_op + # Deduce TMA copy atom and TMA tensor for A, scale, and B + smem_layout_a_per_stage = cute.slice_(self.smem_layout_a, (None, None, None, 0)) + tma_atom_a, tma_tensor_a = cute.nvgpu.make_tiled_tma_atom_A( + a_op, + a, + smem_layout_a_per_stage, + self.mma_tiler, + tiled_mma, + self.cluster_layout_vmnk.shape, + internal_type=( + cutlass.TFloat32 if a.element_type is cutlass.Float32 else None + ), + ) + + tma_atom_scale, tma_tensor_scale = None, None + if cutlass.const_expr(self.scale_mode == TransformMode.ConvertScale): + # Partition smem layout for scale tensor to make it compatible with TMA atom + smem_layout_for_tma_atom = cute.get( + tiled_mma._thrfrg_A(self.smem_layout_scale_per_stage.outer), mode=[1] + ) + # ((MMA_M, MMA_K), REST_M, REST_K) + smem_layout_for_tma_atom = cute.dice( + smem_layout_for_tma_atom, + (1, (1,) * cute.rank(self.smem_layout_scale_per_stage.outer)), + ) + tma_atom_scale, tma_tensor_scale = cute.nvgpu.make_tiled_tma_atom_A( + a_scale_op, + cute.make_tensor(a_scale.iterator, self.gmem_layout_scale), + smem_layout_for_tma_atom, + # (SCALE_M, 1, SCALE_K) + (self.scale_tile_shape[0], 1, self.scale_tile_shape[1]), + tiled_mma, + self.cluster_layout_vmnk.shape, + internal_type=( + cutlass.TFloat32 + if a_scale.element_type is cutlass.Float32 + else None + ), + ) + + smem_layout_b_per_stage = cute.slice_(self.smem_layout_b, (None, None, None, 0)) + tma_atom_b, tma_tensor_b = cute.nvgpu.make_tiled_tma_atom_B( + b_op, + b, + smem_layout_b_per_stage, + self.mma_tiler, + tiled_mma, + self.cluster_layout_vmnk.shape, + internal_type=( + cutlass.TFloat32 if b.element_type is cutlass.Float32 else None + ), + ) + + # Calculate copy size for tensor A, B, and scale + a_copy_size = cute.size_in_bytes(self.a_dtype, smem_layout_a_per_stage) + b_copy_size = cute.size_in_bytes(self.b_dtype, smem_layout_b_per_stage) + a_scale_copy_size = ( + cute.size_in_bytes(self.a_scale_dtype, self.smem_layout_scale_per_stage) + if self.scale_mode is TransformMode.ConvertScale + else 0 + ) + + self.num_tma_load_bytes_a = a_copy_size + self.num_tma_load_bytes_b = b_copy_size * cute.size(tiled_mma.thr_id.shape) + self.num_tma_load_bytes_scale = a_scale_copy_size + self.tile_sched_params, grid = self._compute_grid( + c, + self.cta_tile_shape_mnk, + self.cluster_shape_mn, + max_active_clusters, + ) + + epi_smem_layout = cute.slice_(self.c_smem_layout_staged, (None, None, 0)) + tma_atom_c, tma_tensor_c = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileS2GOp(), + c, + epi_smem_layout, + self.epi_tile, + ) + + @cute.struct + class SharedStorage: + # buffer holding group search results + tile_info: cute.struct.MemRange[cutlass.Int32, 4 * self.num_tile_info_stage] + a_load2trans_full_mbar_ptr: cute.struct.MemRange[ + cutlass.Int64, self.num_load2trans_stage + ] + a_load2trans_empty_mbar_ptr: cute.struct.MemRange[ + cutlass.Int64, self.num_load2trans_stage + ] + a_scale_load2trans_full_mbar_ptr: cute.struct.MemRange[ + cutlass.Int64, self.num_scale_load2trans_stage + ] + a_scale_load2trans_empty_mbar_ptr: cute.struct.MemRange[ + cutlass.Int64, self.num_scale_load2trans_stage + ] + a_trans2mma_full_mbar_ptr: cute.struct.MemRange[ + cutlass.Int64, self.num_trans2mma_stage + ] + a_trans2mma_empty_mbar_ptr: cute.struct.MemRange[ + cutlass.Int64, self.num_trans2mma_stage + ] + b_load2mma_full_mbar_ptr: cute.struct.MemRange[ + cutlass.Int64, self.num_load2trans_stage + ] + b_load2mma_empty_mbar_ptr: cute.struct.MemRange[ + cutlass.Int64, self.num_load2trans_stage + ] + acc_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage] + acc_empty_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage] + tile_info_full_mbar_ptr: cute.struct.MemRange[ + cutlass.Int64, self.num_tile_info_stage + ] + tile_info_empty_mbar_ptr: cute.struct.MemRange[ + cutlass.Int64, self.num_tile_info_stage + ] + tmem_dealloc_mbar_ptr: cutlass.Int64 + tmem_holding_buf: cutlass.Int32 + + self.shared_storage = SharedStorage + + # Launch kernel + self.kernel( + tiled_mma, + tma_atom_a, + tma_tensor_a, + tma_atom_scale, + tma_tensor_scale, + tma_atom_b, + tma_tensor_b, + tma_atom_c, + tma_tensor_c, + c, + cumsum, + self.group_count, + self.cluster_layout_vmnk, + self.smem_layout_a, + self.smem_layout_scale, + self.smem_layout_a_transform, + self.smem_layout_b, + self.c_smem_layout_staged, + self.epi_tile, + self.tile_sched_params, + ).launch( + grid=grid, + block=[self.threads_per_cta, 1, 1], + cluster=(*self.cluster_shape_mn, 1), + min_blocks_per_mp=1, + stream=stream, + ) + return + + # GPU device kernel + @cute.kernel + def kernel( + self, + tiled_mma: cute.TiledMma, + tma_atom_a: cute.CopyAtom, + mA_mkl: cute.Tensor, + tma_atom_s: Optional[cute.CopyAtom], + mS_mkl: Optional[cute.Tensor], + tma_atom_b: cute.CopyAtom, + mB_nkl: cute.Tensor, + tma_atom_c: cute.CopyAtom, + mC_mnl: cute.Tensor, + tensor_c: cute.Tensor, + cumsum: cute.Tensor, + group_count: cutlass.Constexpr[int], + cluster_layout_vmnk: cute.Layout, + a_smem_layout: cute.ComposedLayout, + scale_smem_layout: cute.ComposedLayout, + a_smem_layout_transform: cute.ComposedLayout, + b_smem_layout: cute.ComposedLayout, + c_smem_layout_staged: cute.ComposedLayout, + epi_tile: cute.Tile, + tile_sched_params: utils.PersistentTileSchedulerParams, + ): + """ + GPU device kernel performing the Persistent Mixed-Input Grouped GEMM computation. + """ + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + tidx, _, _ = cute.arch.thread_idx() + bidx, bidy, bidz = cute.arch.block_idx() + # Prefetch TMA descriptors + if warp_idx == self.epilog_warp_id[0]: + cpasync.prefetch_descriptor(tma_atom_a) + cpasync.prefetch_descriptor(tma_atom_b) + if cutlass.const_expr(self.scale_mode == TransformMode.ConvertScale): + cpasync.prefetch_descriptor(tma_atom_s) + cpasync.prefetch_descriptor(tma_atom_c) + + use_2cta_instrs = cute.size(tiled_mma.thr_id.shape) == 2 + bidx, bidy, bidz = cute.arch.block_idx() + # Compute how many k_tiles share the same scale + num_k_tiles_per_scale = self.scale_granularity_k // self.cta_tile_shape_mnk[2] + + mma_tile_coord_v = bidx % cute.size(tiled_mma.thr_id.shape) + is_leader_cta = mma_tile_coord_v == 0 + cta_rank_in_cluster = cute.arch.make_warp_uniform( + cute.arch.block_idx_in_cluster() + ) + block_in_cluster_coord_vmnk = cluster_layout_vmnk.get_flat_coord( + cta_rank_in_cluster + ) + tidx, _, _ = cute.arch.thread_idx() + + smem = utils.SmemAllocator() + storage = smem.allocate(self.shared_storage) + + # Initialize load2transform pipeline, which tracks the dependencies between TMA's loading + # of A and B, and the transformation of A and MMA's consumption + transform_thread_idx = ( + tidx - 32 * self.transform_warp_id[0] + if tidx >= 32 * self.transform_warp_id[0] + else tidx + ) + a_load2trans_pipeline = pipeline.PipelineTmaAsync.create( + barrier_storage=storage.a_load2trans_full_mbar_ptr.data_ptr(), + num_stages=self.num_load2trans_stage, + producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), + consumer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, + self.num_mcast_ctas_a * len(self.transform_warp_id), + ), + tx_count=self.num_tma_load_bytes_a, + cta_layout_vmnk=cluster_layout_vmnk, + tidx=transform_thread_idx, + mcast_mode_mn=(1, 0), # multicast for A will only happen on the M-mode + defer_sync=True, + ) + # Initialize scale_load2trans pipeline, which tracks the dependencies between TMA's loading + # of scale, and the transformation of A + scale_load2trans_pipeline = None + if cutlass.const_expr(self.scale_mode == TransformMode.ConvertScale): + num_producers_a_scale = self.num_mcast_ctas_a + scale_load2trans_pipeline = pipeline.PipelineTmaAsync.create( + barrier_storage=storage.a_scale_load2trans_full_mbar_ptr.data_ptr(), + num_stages=self.num_scale_load2trans_stage, + producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), + consumer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, + num_producers_a_scale + * len(self.transform_warp_id) + * num_k_tiles_per_scale, + ), + tx_count=self.num_tma_load_bytes_scale, + cta_layout_vmnk=cluster_layout_vmnk, + tidx=transform_thread_idx, + mcast_mode_mn=( + 1, + 0, + ), # multicast for scale_a will only happen on the M-mode + defer_sync=True, + ) + # Initialize transform2mma pipeline, which tracks the dependencies between the transformation + # of A and MMA's consumption of transformed A + cta_v_size = cute.size(cluster_layout_vmnk, mode=[0]) + trans2mma_pipeline = pipeline.PipelineAsyncUmma.create( + barrier_storage=storage.a_trans2mma_full_mbar_ptr.data_ptr(), + num_stages=self.num_trans2mma_stage, + producer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, + 32 * len(self.transform_warp_id) * cta_v_size, + ), + consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ) + # Initialize pipeline for tensor B load to MMA + # MMA warp informs TMA warp to proceed to load next tile of B tensor + b_load2mma_pipeline = pipeline.PipelineTmaUmma.create( + barrier_storage=storage.b_load2mma_full_mbar_ptr.data_ptr(), + num_stages=self.num_load2trans_stage, + producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), + consumer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, self.num_mcast_ctas_b + ), + tx_count=self.num_tma_load_bytes_b, + cta_layout_vmnk=cluster_layout_vmnk, + mcast_mode_mn=(0, 1), # multicast for B will only happen on the N-mode + defer_sync=True, + ) + # Initialize accumulator pipeline, which tracks the dependencies between + # MMA's computation of accumulators and epilogue warps' consumption of accumulators + acc_pipeline = pipeline.PipelineUmmaAsync.create( + barrier_storage=storage.acc_full_mbar_ptr.data_ptr(), + num_stages=self.num_acc_stage, + producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), + consumer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, cta_v_size * len(self.epilog_warp_id) + ), + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ) + # Initialize tile info pipeline, which tracks the dependencies between + # tile scheduling warp and other warps + # Skip scheduler warp and TMA scale load warp when scale_mode is ConvertOnly + # when computing consumer thread count + num_tile_info_pipeline_consumer_threads = ( + self.threads_per_cta + - 32 + - (32 if self.scale_mode is TransformMode.ConvertOnly else 0) + ) + tile_info_pipeline = pipeline.PipelineAsync.create( + barrier_storage=storage.tile_info_full_mbar_ptr.data_ptr(), + num_stages=self.num_tile_info_stage, + producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * 1), + consumer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, + num_tile_info_pipeline_consumer_threads, + ), + defer_sync=True, + ) + + # Tensor memory dealloc barrier init + tmem = utils.TmemAllocator( + storage.tmem_holding_buf, + barrier_for_retrieve=self.tmem_ptr_sync_barrier, + allocator_warp_id=self.epilog_warp_id[0], + is_two_cta=use_2cta_instrs, + two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar_ptr, + ) + + # Cluster arrive after barrier init + pipeline_init_arrive(cluster_shape_mn=self.cluster_shape_mn, is_relaxed=True) + + # Setup smem tensor A/scale/B/C + sC = smem.allocate_tensor( + element_type=self.c_dtype, + layout=c_smem_layout_staged.outer, + byte_alignment=self.smem_buffer_align_bytes, + swizzle=c_smem_layout_staged.inner, + ) + sA_input = smem.allocate_tensor( + element_type=self.a_dtype, + layout=a_smem_layout.outer, + byte_alignment=self.smem_buffer_align_bytes, + swizzle=a_smem_layout.inner, + ) + sS_input = ( + smem.allocate_tensor( + element_type=self.mma_dtype, + layout=scale_smem_layout.outer, + byte_alignment=self.smem_buffer_align_bytes, + swizzle=scale_smem_layout.inner, + ) + if self.scale_mode is TransformMode.ConvertScale + else None + ) + sB_input = smem.allocate_tensor( + element_type=self.b_dtype, + layout=b_smem_layout.outer, + byte_alignment=self.smem_buffer_align_bytes, + swizzle=b_smem_layout.inner, + ) + sA_transform = None + # Get smem tensor for transformed A when transform_a_source is SMEM + if cutlass.const_expr(self.transform_a_source == tcgen05.OperandSource.SMEM): + sA_transform = smem.allocate_tensor( + element_type=self.mma_dtype, + layout=a_smem_layout_transform.outer, + byte_alignment=self.smem_buffer_align_bytes, + swizzle=a_smem_layout_transform.inner, + ) + sTile_info = storage.tile_info.get_tensor( + cute.make_layout((4, self.num_tile_info_stage), stride=(1, 4)) + ) + + # Compute multicast mask for A/B buffer full + a_full_mcast_mask = None + b_full_mcast_mask = None + s_full_mcast_mask = None + if cutlass.const_expr(self.is_a_mcast or self.is_b_mcast or use_2cta_instrs): + a_full_mcast_mask = cpasync.create_tma_multicast_mask( + cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=2 + ) + # Scale tensor shares the same multicast mask as the A tensor + s_full_mcast_mask = a_full_mcast_mask + b_full_mcast_mask = cpasync.create_tma_multicast_mask( + cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=1 + ) + + # local_tile partition global tensors + # (bM, bK, loopM, loopK, loopL) + gA_mkl = cute.local_tile( + mA_mkl, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None) + ) + # (bM, bK, loopM, loopK, loopL) + gS_mkl = ( + cute.local_tile( + mS_mkl, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None) + ) + if self.scale_mode is TransformMode.ConvertScale + else None + ) + # (bN, bK, loopN, loopK, loopL) + gB_nkl = cute.local_tile( + mB_nkl, cute.slice_(self.mma_tiler, (0, None, None)), (None, None, None) + ) + # (bM, bN, loopM, loopN, loopL) + gC_mnl = cute.local_tile( + mC_mnl, cute.slice_(self.mma_tiler, (None, None, 0)), (None, None, None) + ) + gC_mnl_simt = cute.local_tile( + tensor_c, cute.slice_(self.mma_tiler, (None, None, 0)), (None, None, None) + ) + k_tile_cnt = cute.size(gA_mkl, mode=[3]) + + # Partition global tensor for TiledMMA_A/B/C + thr_mma = tiled_mma.get_slice(mma_tile_coord_v) + # (MMA, MMA_M, MMA_K, loopM, loopK, loopL) + tCgA = thr_mma.partition_A(gA_mkl) + # (MMA, MMA_M, MMA_K, loopM, loopK, loopL) + tCgS = ( + thr_mma.partition_A(gS_mkl) + if self.scale_mode is TransformMode.ConvertScale + else None + ) + # (MMA, MMA_N, MMA_K, loopN, loopK, loopL) + tCgB = thr_mma.partition_B(gB_nkl) + # (MMA, MMA_M, MMA_N, loopM, loopN, loopL) + tCgC = thr_mma.partition_C(gC_mnl) + tCgC_simt = thr_mma.partition_C(gC_mnl_simt) + + # Setup copy atom to load A from shared memory for further transformation + copy_atom_a_input = ( + cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), self.a_dtype, num_bits_per_copy=32 + ) + if self.scale_mode is TransformMode.ConvertScale + else None + ) + a_smem_shape = tiled_mma.partition_shape_A( + cute.dice(self.mma_tiler, (1, None, 1)) + ) + # Setup copy atom to store transformed A into tensor memory or shared memory + copy_atom_a_transform = mixed_input_utils.get_copy_atom_a_transform( + self.mma_dtype, + self.use_2cta_instrs, + self.transform_a_source, + a_smem_shape, + self.a_dtype, + ) + + # Partition global/shared tensor for TMA load A/B + # TMA load A partition_S/D + a_cta_layout = cute.make_layout( + cute.slice_(cluster_layout_vmnk, (0, 0, None, 0)).shape + ) + # ((atom_v, rest_v), STAGE) + # ((atom_v, rest_v), loopM, loopK, loopL) + tAsA, tAgA = cpasync.tma_partition( + tma_atom_a, + block_in_cluster_coord_vmnk[2], + a_cta_layout, + cute.group_modes(sA_input, 0, 3), + cute.group_modes(tCgA, 0, 3), + ) + + tCsS = None + tSsS = None + tSgS = None + if cutlass.const_expr(self.scale_mode == TransformMode.ConvertScale): + thr_mma_leader_cta = tiled_mma.get_slice(0) + # (MMA, MMA_M, MMA_K, STAGE) + tCsS = thr_mma_leader_cta.partition_A(sS_input) + # ((atom_v, rest_v), STAGE) + # ((atom_v, rest_v), loopM, loopK, loopL) + tSsS, tSgS = mixed_input_utils.scale_tma_partition( + tCsS, + tCgS, + tma_atom_s, + block_in_cluster_coord_vmnk, + a_cta_layout, + ) + + # TMA load B partition_S/D + b_cta_layout = cute.make_layout( + cute.slice_(cluster_layout_vmnk, (0, None, 0, 0)).shape + ) + # ((atom_v, rest_v), STAGE) + # ((atom_v, rest_v), loopM, loopK, loopL) + tBsB, tBgB = cpasync.tma_partition( + tma_atom_b, + block_in_cluster_coord_vmnk[1], + b_cta_layout, + cute.group_modes(sB_input, 0, 3), + cute.group_modes(tCgB, 0, 3), + ) + + # (MMA, MMA_N, MMA_K, STAGE) + tCrB = tiled_mma.make_fragment_B(sB_input) + # (MMA, MMA_M, MMA_N) + acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2]) + tCtAcc_fake = tiled_mma.make_fragment_C( + cute.append(acc_shape, self.num_acc_stage) + ) + + # Cluster wait before TMEM alloc and ensure pipelines are ready + pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mn) + + # TMEM allocation + tmem.allocate(self.num_tmem_alloc_cols) + tmem.wait_for_alloc() + + # Schedule warp + if warp_idx == self.schedule_warp_id: + cute.arch.setmaxregister_decrease(self.num_regs_schedule_warp) + # Persistent tile scheduling loop + tile_sched = utils.StaticPersistentRuntimeTileScheduler.create( + tile_sched_params, + (bidx, bidy, bidz), + cute.arch.grid_dim(), + inner_mode=0, + ) + work_tile = tile_sched.initial_work_tile_info() + tile_info_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.num_tile_info_stage + ) + # Create initial group search state + search_state = ( + mixed_input_utils.create_initial_contiguous_group_search_state() + ) + not_last_tile = cutlass.Boolean(1) + while not_last_tile: + tile_info_pipeline.producer_acquire(tile_info_producer_state) + cluster_tile_coord_mnl = work_tile.tile_idx + cta_tile_coord_m = ( + cluster_tile_coord_mnl[0] * self.cluster_shape_mn[0] + + block_in_cluster_coord_vmnk[1] * cute.size(tiled_mma.thr_id.shape) + + block_in_cluster_coord_vmnk[0] + ) + cta_tile_offset_n = block_in_cluster_coord_vmnk[2] + search_state = mixed_input_utils.contiguous_group_search( + self.cluster_tile_shape_mnk, + group_count, + cluster_tile_coord_mnl[1], + search_state, + cumsum, + 1, # mode index to perform the search. 0 for M and 1 for N + ) + cur_sTile_info = sTile_info[(None, tile_info_producer_state.index)] + not_last_tile = search_state.cur_group_idx <= group_count + # Store tile info into shared memory buffer + with cute.arch.elect_one(): + cur_sTile_info[0] = cta_tile_coord_m + cur_sTile_info[1] = ( + search_state.cur_start + + cta_tile_offset_n * self.cta_tile_shape_mnk[1] + ) + cur_sTile_info[2] = search_state.cur_group_idx - 1 + cur_sTile_info[3] = ( + search_state.cur_boundary + - search_state.cur_start + - (cta_tile_offset_n * self.cta_tile_shape_mnk[1]) + ) + # Fence and barrier to ensure tile info store has finished + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) + self.sched_sync_barrier.arrive_and_wait() + # Commit tile info pipeline + tile_info_pipeline.producer_commit(tile_info_producer_state) + # Advance to next tile + tile_info_producer_state.advance() + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + tile_info_pipeline.producer_tail(tile_info_producer_state) + + # Specialized TMA load warp for A/B tensor + if warp_idx == self.tma_warp_id: + cute.arch.setmaxregister_decrease(self.num_regs_tma_warps) + # Persistent tile scheduling loop + tile_info_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_tile_info_stage + ) + tile_info_pipeline.consumer_wait(tile_info_consumer_state) + work_tile = mixed_input_utils.make_contiguous_group_work_tile_info( + group_count, sTile_info[(None, tile_info_consumer_state.index)] + ) + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) + tile_info_pipeline.consumer_release(tile_info_consumer_state) + tile_info_consumer_state.advance() + a_load2trans_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.num_load2trans_stage + ) + b_load2mma_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.num_load2trans_stage + ) + + while work_tile.is_valid_tile: + tAgA_slice = tAgA[ + ( + None, + work_tile.cta_coord_m // cute.size(tiled_mma.thr_id.shape), + None, + work_tile.group_idx, + ) + ] + # Apply offset to B tensor based on group search result + coord_n_offset = ( + (work_tile.coord_n, 0, 0) + if cutlass.const_expr( + self.b_major_mode == tcgen05.OperandMajorMode.MN + ) + else (0, work_tile.coord_n, 0) + ) + tBgB_slice = cute.make_tensor( + ( + tBgB.iterator[0] + coord_n_offset[0], + coord_n_offset[1] + tBgB.iterator[1], + coord_n_offset[2] + tBgB.iterator[2], + ), + cute.slice_(tBgB.layout, (None, 0, None, 0)), + ) + + a_load2trans_producer_state.reset_count() + peek_load2trans_empty_status = cutlass.Boolean(1) + if a_load2trans_producer_state.count < k_tile_cnt: + peek_load2trans_empty_status = ( + a_load2trans_pipeline.producer_try_acquire( + a_load2trans_producer_state + ) + ) + b_load2mma_producer_state.reset_count() + for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1): + a_load2trans_pipeline.producer_acquire( + a_load2trans_producer_state, peek_load2trans_empty_status + ) + b_load2mma_pipeline.producer_acquire(b_load2mma_producer_state) + # TMA load A/B + cute.copy( + tma_atom_a, + tAgA_slice[(None, a_load2trans_producer_state.count)], + tAsA[(None, a_load2trans_producer_state.index)], + tma_bar_ptr=a_load2trans_pipeline.producer_get_barrier( + a_load2trans_producer_state + ), + mcast_mask=a_full_mcast_mask, + ) + cute.copy( + tma_atom_b, + tBgB_slice[(None, b_load2mma_producer_state.count)], + tBsB[(None, b_load2mma_producer_state.index)], + tma_bar_ptr=b_load2mma_pipeline.producer_get_barrier( + b_load2mma_producer_state + ), + mcast_mask=b_full_mcast_mask, + ) + a_load2trans_pipeline.producer_commit(a_load2trans_producer_state) + b_load2mma_pipeline.producer_commit(b_load2mma_producer_state) + a_load2trans_producer_state.advance() + b_load2mma_producer_state.advance() + if a_load2trans_producer_state.count < k_tile_cnt: + peek_load2trans_empty_status = ( + a_load2trans_pipeline.producer_try_acquire( + a_load2trans_producer_state + ) + ) + # Advance to next tile + tile_info_pipeline.consumer_wait(tile_info_consumer_state) + work_tile = mixed_input_utils.make_contiguous_group_work_tile_info( + group_count, sTile_info[(None, tile_info_consumer_state.index)] + ) + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) + tile_info_pipeline.consumer_release(tile_info_consumer_state) + tile_info_consumer_state.advance() + # Wait A/B buffer empty + a_load2trans_pipeline.producer_tail(a_load2trans_producer_state) + b_load2mma_pipeline.producer_tail(b_load2mma_producer_state) + + # Specialized TMA load for scale tensor + if warp_idx == self.scale_tma_warp_id: + cute.arch.setmaxregister_decrease(self.num_regs_tma_warps) + if cutlass.const_expr(self.scale_mode == TransformMode.ConvertScale): + # Persistent tile scheduling loop + tile_info_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_tile_info_stage + ) + tile_info_pipeline.consumer_wait(tile_info_consumer_state) + work_tile = mixed_input_utils.make_contiguous_group_work_tile_info( + group_count, sTile_info[(None, tile_info_consumer_state.index)] + ) + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) + tile_info_pipeline.consumer_release(tile_info_consumer_state) + tile_info_consumer_state.advance() + scale_load2trans_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.num_scale_load2trans_stage + ) + scale_k_tile_cnt = cute.size(mS_mkl.layout.shape[1][1]) + + while work_tile.is_valid_tile: + # ((atom_v, rest_v), RestK) + tSgS_slice = tSgS[ + ( + None, + work_tile.cta_coord_m // cute.size(tiled_mma.thr_id.shape), + None, + work_tile.group_idx, + ) + ] + # Filter zeros in rest mode + rest_filtered = cute.filter_zeros(tSgS_slice[(0, None)].layout) + tSgS_slice_filtered = cute.make_tensor( + tSgS_slice.iterator, + cute.make_layout( + (tSgS_slice.layout[0].shape, rest_filtered.shape), + stride=(tSgS_slice.layout[0].stride, rest_filtered.stride), + ), + ) + + scale_load2trans_producer_state.reset_count() + peek_scale_load2trans_empty_status = cutlass.Boolean(1) + if scale_load2trans_producer_state.count < scale_k_tile_cnt: + peek_scale_load2trans_empty_status = ( + scale_load2trans_pipeline.producer_try_acquire( + scale_load2trans_producer_state + ) + ) + for k_tile in cutlass.range(0, scale_k_tile_cnt, 1, unroll=1): + scale_load2trans_pipeline.producer_acquire( + scale_load2trans_producer_state, + peek_scale_load2trans_empty_status, + ) + # TMA load scale + cute.copy( + tma_atom_s, + tSgS_slice_filtered[ + (None, scale_load2trans_producer_state.count) + ], + tSsS[(None, scale_load2trans_producer_state.index)], + tma_bar_ptr=scale_load2trans_pipeline.producer_get_barrier( + scale_load2trans_producer_state + ), + mcast_mask=s_full_mcast_mask, + ) + + scale_load2trans_producer_state.advance() + peek_scale_load2trans_empty_status = cutlass.Boolean(1) + if scale_load2trans_producer_state.count < scale_k_tile_cnt: + peek_scale_load2trans_empty_status = ( + scale_load2trans_pipeline.producer_try_acquire( + scale_load2trans_producer_state + ) + ) + # Advance to next tile + tile_info_pipeline.consumer_wait(tile_info_consumer_state) + work_tile = mixed_input_utils.make_contiguous_group_work_tile_info( + group_count, sTile_info[(None, tile_info_consumer_state.index)] + ) + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) + tile_info_pipeline.consumer_release(tile_info_consumer_state) + tile_info_consumer_state.advance() + # Wait scale buffer empty + scale_load2trans_pipeline.producer_tail(scale_load2trans_producer_state) + + # Specialized transform warps + if warp_idx >= self.transform_warp_id[0]: + cute.arch.setmaxregister_increase(self.num_regs_transform_warps) + transform_local_tidx = tidx - 32 * self.transform_warp_id[0] + # Get the pointer to the TMEM buffer + tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + accumulators = cute.make_tensor(tmem_ptr, tCtAcc_fake.layout) + + tCrA = None + if cutlass.const_expr(self.transform_a_source == tcgen05.OperandSource.TMEM): + tmem_ptr_transform = cute.recast_ptr( + accumulators.iterator + self.num_acc_tmem_cols, dtype=self.mma_dtype + ) + tCrA = cute.make_tensor( + tmem_ptr_transform, + tiled_mma.make_fragment_A(a_smem_layout_transform.outer).layout, + ) + else: + tCrA = tiled_mma.make_fragment_A(sA_transform) + # Partition tensors for transform input and output and set up the copy atom + # used for loading and storing transformed A tensor + src_copy_a, dst_copy_a, tAsA_input, tAsA_transform = ( + mixed_input_utils.transform_partition( + self.transform_a_source, + self.scale_mode, + copy_atom_a_input, + copy_atom_a_transform, + sA_input, + ( + tCrA + if self.transform_a_source == tcgen05.OperandSource.TMEM + else sA_transform + ), + transform_local_tidx, + ) + ) + # make fragment for input A and transformed A + tArA = cute.make_rmem_tensor( + tAsA_input[(None, None, None, None, 0)].shape, tAsA_input.element_type + ) + tArA_transform = cute.make_rmem_tensor( + tAsA_input[(None, None, None, None, 0)].shape, self.mma_dtype + ) + # Partition scale tensor + smem_thr_copy_S = None + tSsS_trans = None + tSrS_copy = None + tSrS = None + if cutlass.const_expr(self.scale_mode == TransformMode.ConvertScale): + smem_thr_copy_S, tSsS_trans, tSrS_copy, tSrS = ( + mixed_input_utils.scale_partition( + src_copy_a, tCsS, transform_local_tidx, self.mma_dtype + ) + ) + assert cute.size(tSrS, mode=[0]) == cute.size(tArA, mode=[0]), ( + "tSrS and tArA have different leading dimension" + ) + assert cute.size(tSrS) == cute.size(tArA), ( + "tSrS and tArA have different shape" + ) + # Deduce a sub-tile size and tile tensors + transform_tiler_size = min( + cute.size(cute.coalesce(tAsA_input.layout), mode=[0]), 64 + ) + transform_tiler = cute.make_layout(transform_tiler_size) + tArA_load = cute.flat_divide(tArA, transform_tiler) + tArA_load = cute.group_modes(tArA_load, 1, cute.rank(tArA_load)) + tSrS_load = ( + cute.flat_divide(tSrS, transform_tiler) + if self.scale_mode is TransformMode.ConvertScale + else None + ) + tSrS_load = ( + cute.group_modes(tSrS_load, 1, cute.rank(tSrS_load)) + if self.scale_mode is TransformMode.ConvertScale + else None + ) + tArA_transform_store = cute.flat_divide(tArA_transform, transform_tiler) + tArA_transform_store = cute.group_modes( + tArA_transform_store, 1, cute.rank(tArA_transform_store) + ) + + tile_info_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_tile_info_stage + ) + tile_info_pipeline.consumer_wait(tile_info_consumer_state) + work_tile = mixed_input_utils.make_contiguous_group_work_tile_info( + group_count, sTile_info[(None, tile_info_consumer_state.index)] + ) + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) + tile_info_pipeline.consumer_release(tile_info_consumer_state) + tile_info_consumer_state.advance() + a_load2trans_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, + self.num_load2trans_stage, + ) + scale_load2trans_consumer_state = ( + pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, + self.num_scale_load2trans_stage, + ) + if self.scale_mode is TransformMode.ConvertScale + else None + ) + trans2mma_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, + self.num_trans2mma_stage, + ) + while work_tile.is_valid_tile: + a_load2trans_consumer_state.reset_count() + peek_load2trans_full_status = cutlass.Boolean(1) + if a_load2trans_consumer_state.count < k_tile_cnt: + peek_load2trans_full_status = ( + a_load2trans_pipeline.consumer_try_wait( + a_load2trans_consumer_state + ) + ) + peek_scale_load2trans_full_status = cutlass.Boolean(1) + if cutlass.const_expr(self.scale_mode == TransformMode.ConvertScale): + scale_load2trans_consumer_state.reset_count() + peek_scale_load2trans_full_status = ( + scale_load2trans_pipeline.consumer_try_wait( + scale_load2trans_consumer_state + ) + ) + trans2mma_producer_state.reset_count() + peek_trans2mma_empty_status = cutlass.Boolean(1) + if trans2mma_producer_state.count < k_tile_cnt: + peek_trans2mma_empty_status = ( + trans2mma_pipeline.producer_try_acquire( + trans2mma_producer_state + ) + ) + + for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1): + a_load2trans_pipeline.consumer_wait( + a_load2trans_consumer_state, peek_load2trans_full_status + ) + tAsA_input_slice = tAsA_input[ + (None, None, None, None, a_load2trans_consumer_state.index) + ] + tAsA_input_slice = cute.flat_divide( + tAsA_input_slice, transform_tiler + ) + tAsA_input_slice = cute.group_modes( + tAsA_input_slice, 1, cute.rank(tAsA_input_slice) + ) + if cutlass.const_expr( + self.scale_mode == TransformMode.ConvertScale + ): + scale_load2trans_pipeline.consumer_wait( + scale_load2trans_consumer_state, + peek_scale_load2trans_full_status, + ) + trans2mma_pipeline.producer_acquire( + trans2mma_producer_state, peek_trans2mma_empty_status + ) + # load scale tensor when needed + if cutlass.const_expr( + self.scale_mode == TransformMode.ConvertScale + ): + if k_tile % num_k_tiles_per_scale == 0: + tSsS_slice = tSsS_trans[ + ( + None, + None, + None, + None, + scale_load2trans_consumer_state.index, + ) + ] + tSsS_slice_filtered = cute.make_tensor( + tSsS_slice.iterator, + cute.filter_zeros(tSsS_slice.layout), + ) + cute.autovec_copy(tSsS_slice_filtered, tSrS_copy) + cur_scale_load2trans_consumer_state = ( + scale_load2trans_consumer_state.clone() + ) + if (k_tile + 1) % num_k_tiles_per_scale == 0: + scale_load2trans_consumer_state.advance() + + cur_a_load2trans_consumer_state = ( + a_load2trans_consumer_state.clone() + ) + for idx in cutlass.range_constexpr(cute.size(tArA_load, mode=[1])): + # Load A from shared memory + cute.autovec_copy( + tAsA_input_slice[(None, idx)], + tArA_load[(None, idx)], + ) + if cutlass.const_expr( + idx == cute.size(tArA_load, mode=[1]) - 1 + ): + a_load2trans_consumer_state.advance() + if a_load2trans_consumer_state.count < k_tile_cnt: + peek_load2trans_full_status = ( + a_load2trans_pipeline.consumer_try_wait( + a_load2trans_consumer_state + ) + ) + if cutlass.const_expr( + self.scale_mode == TransformMode.ConvertScale + ): + peek_scale_load2trans_full_status = ( + scale_load2trans_pipeline.consumer_try_wait( + scale_load2trans_consumer_state + ) + ) + # Convert it to mma dtype + tensor_transformed = mixed_input_utils.cvt_tensor_a( + tArA_load[(None, idx)], self.mma_dtype, self.shuffle_a + ) + if cutlass.const_expr( + self.scale_mode == TransformMode.ConvertScale + ): + scale = cute.TensorSSA( + tSrS_load[(None, idx)].load(), + tensor_transformed.shape, + self.mma_dtype, + ) + # Apply scale + tensor_transformed = tensor_transformed * scale + tArA_transform_store[(None, idx)].store(tensor_transformed) + # Store transformed A to tensor memory or shared memory + mixed_input_utils.store_transformed_a( + tArA_transform, + tAsA_transform[ + (None, None, None, None, trans2mma_producer_state.index) + ], + dst_copy_a, + ) + # Ensure all transform threads have finished the copy and reached the fence + self.transform_sync_barrier.arrive_and_wait() + if cutlass.const_expr( + self.transform_a_source == tcgen05.OperandSource.TMEM + ): + cute.arch.fence_view_async_tmem_store() + else: + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) + if cutlass.const_expr( + self.scale_mode == TransformMode.ConvertScale + ): + scale_load2trans_pipeline.consumer_release( + cur_scale_load2trans_consumer_state + ) + + a_load2trans_pipeline.consumer_release( + cur_a_load2trans_consumer_state + ) + # Signal the completion of transformation + trans2mma_pipeline.producer_commit(trans2mma_producer_state) + trans2mma_producer_state.advance() + if trans2mma_producer_state.count < k_tile_cnt: + peek_trans2mma_empty_status = ( + trans2mma_pipeline.producer_try_acquire( + trans2mma_producer_state + ) + ) + # Advance to next tile + tile_info_pipeline.consumer_wait(tile_info_consumer_state) + work_tile = mixed_input_utils.make_contiguous_group_work_tile_info( + group_count, sTile_info[(None, tile_info_consumer_state.index)] + ) + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) + tile_info_pipeline.consumer_release(tile_info_consumer_state) + tile_info_consumer_state.advance() + # Wait a_transform buffer empty + trans2mma_pipeline.producer_tail(trans2mma_producer_state) + + # Specialized MMA warp + if warp_idx == self.mma_warp_id: + cute.arch.setmaxregister_decrease(self.num_regs_mma_warp) + # Get the pointer to the TMEM buffer + tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + accumulators = cute.make_tensor(tmem_ptr, tCtAcc_fake.layout) + tCrA = None + if cutlass.const_expr(self.transform_a_source == tcgen05.OperandSource.TMEM): + tmem_ptr_transform = cute.recast_ptr( + accumulators.iterator + self.num_acc_tmem_cols, dtype=self.mma_dtype + ) + tCrA = cute.make_tensor( + tmem_ptr_transform, + tiled_mma.make_fragment_A(a_smem_layout_transform.outer).layout, + ) + else: + tCrA = tiled_mma.make_fragment_A(sA_transform) + tCtAcc_base = accumulators + # Persistent tile scheduling loop + tile_info_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_tile_info_stage + ) + tile_info_pipeline.consumer_wait(tile_info_consumer_state) + work_tile = mixed_input_utils.make_contiguous_group_work_tile_info( + group_count, sTile_info[(None, tile_info_consumer_state.index)] + ) + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) + tile_info_pipeline.consumer_release(tile_info_consumer_state) + tile_info_consumer_state.advance() + trans2mma_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_trans2mma_stage + ) + b_load2mma_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_load2trans_stage + ) + acc_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.num_acc_stage + ) + while work_tile.is_valid_tile: + # (MMA, MMA_M, MMA_N) + tCtAcc = tCtAcc_base[(None, None, None, acc_producer_state.index)] + b_load2mma_consumer_state.reset_count() + trans2mma_consumer_state.reset_count() + peek_trans2mma_full_status = cutlass.Boolean(1) + if is_leader_cta: + if trans2mma_consumer_state.count < k_tile_cnt: + peek_trans2mma_full_status = ( + trans2mma_pipeline.consumer_try_wait( + trans2mma_consumer_state + ) + ) + acc_pipeline.producer_acquire(acc_producer_state) + + tiled_mma.set(tcgen05.Field.ACCUMULATE, False) + # Mma mainloop + for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1): + trans2mma_pipeline.consumer_wait( + trans2mma_consumer_state, peek_trans2mma_full_status + ) + b_load2mma_pipeline.consumer_wait(b_load2mma_consumer_state) + num_kblocks = cute.size(tCrA, mode=[2]) + for kblock_idx in cutlass.range(num_kblocks, unroll_full=True): + kblock_coord_a = ( + None, + None, + kblock_idx, + trans2mma_consumer_state.index, + ) + kblock_coord_b = ( + None, + None, + kblock_idx, + b_load2mma_consumer_state.index, + ) + + cute.gemm( + tiled_mma, + tCtAcc, + tCrA[kblock_coord_a], + tCrB[kblock_coord_b], + tCtAcc, + ) + # Enable accumulate on tCtAcc after first kblock + tiled_mma.set(tcgen05.Field.ACCUMULATE, True) + trans2mma_pipeline.consumer_release(trans2mma_consumer_state) + b_load2mma_pipeline.consumer_release(b_load2mma_consumer_state) + trans2mma_consumer_state.advance() + b_load2mma_consumer_state.advance() + peek_trans2mma_full_status = cutlass.Boolean(1) + if trans2mma_consumer_state.count < k_tile_cnt: + peek_trans2mma_full_status = ( + trans2mma_pipeline.consumer_try_wait( + trans2mma_consumer_state + ) + ) + # Async arrive accumulator buffer full + acc_pipeline.producer_commit(acc_producer_state) + acc_producer_state.advance() + + # Advance to next tile + tile_info_pipeline.consumer_wait(tile_info_consumer_state) + work_tile = mixed_input_utils.make_contiguous_group_work_tile_info( + group_count, sTile_info[(None, tile_info_consumer_state.index)] + ) + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) + tile_info_pipeline.consumer_release(tile_info_consumer_state) + tile_info_consumer_state.advance() + # Wait for accumulator buffer empty + acc_pipeline.producer_tail(acc_producer_state) + + # Specialized epilogue warps + if warp_idx < self.mma_warp_id: + cute.arch.setmaxregister_increase(self.num_regs_epilogue_warps) + epi_tidx = tidx + # Get the pointer to the TMEM buffer + tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + accumulators = cute.make_tensor(tmem_ptr, tCtAcc_fake.layout) + tCtAcc_base = accumulators + # Partition for epilogue + tiled_copy_t2r, tTR_tAcc_base, tTR_rAcc = ( + mixed_input_utils.epilog_tmem_copy_and_partition( + self.cta_tile_shape_mnk, + self.c_layout, + self.c_dtype, + self.acc_dtype, + epi_tidx, + tCtAcc_base, + tCgC, + epi_tile, + self.use_2cta_instrs, + ) + ) + + tTR_rC = cute.make_rmem_tensor(tTR_rAcc.shape, self.c_dtype) + tiled_copy_r2s, tRS_rC, tRS_sC = ( + mixed_input_utils.epilog_smem_copy_and_partition( + self.c_layout, + self.c_dtype, + self.acc_dtype, + tiled_copy_t2r, + tTR_rC, + epi_tidx, + sC, + ) + ) + (tma_atom_c, bSG_sC, bSG_gC_partitioned, simt_atom, tTR_gC_partitioned) = ( + mixed_input_utils.epilog_gmem_copy_and_partition( + self.c_dtype, + epi_tidx, + tma_atom_c, + tiled_copy_t2r, + tCgC, + tCgC_simt, + epi_tile, + sC, + ) + ) + + # Predicates + thr_mapping = cute.make_identity_tensor( + (self.cta_tile_shape_mnk[0], self.cta_tile_shape_mnk[1]) + ) + thr_mapping_mn = cute.flat_divide(thr_mapping, epi_tile) + thr_copy_t2r = tiled_copy_t2r.get_slice(epi_tidx) + m_thr_offset = thr_copy_t2r.partition_D(thr_mapping_mn) + m_thr_offset = cute.group_modes(m_thr_offset, 3, cute.rank(m_thr_offset)) + + acc_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_acc_stage + ) + + c_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + 32 * len(self.epilog_warp_id), + ) + c_pipeline = pipeline.PipelineTmaStore.create( + num_stages=self.num_c_stage, + producer_group=c_producer_group, + ) + + # Persistent tile scheduling loop + tile_info_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_tile_info_stage + ) + tile_info_pipeline.consumer_wait(tile_info_consumer_state) + work_tile = mixed_input_utils.make_contiguous_group_work_tile_info( + group_count, sTile_info[(None, tile_info_consumer_state.index)] + ) + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) + tile_info_pipeline.consumer_release(tile_info_consumer_state) + tile_info_consumer_state.advance() + num_prev_subtiles = cutlass.Int32(0) + while work_tile.is_valid_tile: + bSG_gC = bSG_gC_partitioned[ + ( + None, + None, + None, + work_tile.cta_coord_m // cute.size(tiled_mma.thr_id.shape), + 0, + 0, + ) + ] + tma_store_offset_coord = ( + (work_tile.coord_n, 0, 0) + if cutlass.const_expr(self.c_layout.is_n_major_c()) + else (0, work_tile.coord_n, 0) + ) + bSG_gC = cute.make_tensor( + ( + tma_store_offset_coord[0] + bSG_gC.iterator[0], + tma_store_offset_coord[1] + bSG_gC.iterator[1], + tma_store_offset_coord[2] + bSG_gC.iterator[2], + ), + bSG_gC.layout, + ) + tTR_gC = tTR_gC_partitioned[ + ( + None, + None, + None, + None, + None, + work_tile.cta_coord_m // cute.size(tiled_mma.thr_id.shape), + 0, + 0, + ) + ] + tTR_gC = cute.make_tensor( + tTR_gC.iterator + (work_tile.coord_n * tensor_c.layout.stride[1]), + tTR_gC.layout, + ) + + tTR_tAcc = tTR_tAcc_base[ + (None, None, None, None, None, acc_consumer_state.index) + ] + # Wait for accumulator buffer full + acc_pipeline.consumer_wait(acc_consumer_state) + + tTR_tAcc = cute.group_modes(tTR_tAcc, 3, cute.rank(tTR_tAcc)) + bSG_gC = cute.group_modes(bSG_gC, 1, cute.rank(bSG_gC)) + tTR_gC = cute.group_modes(tTR_gC, 3, cute.rank(tTR_gC)) + + # Store accumulator to global memory in subtiles + subtile_cnt = cute.size(tTR_tAcc.shape, mode=[3]) + for subtile_idx in cutlass.range(subtile_cnt): + # Load accumulator from tensor memory buffer to register + tTR_tAcc_mn = tTR_tAcc[(None, None, None, subtile_idx)] + cute.copy(tiled_copy_t2r, tTR_tAcc_mn, tTR_rAcc) + if work_tile.distance_to_boundary >= self.cta_tile_shape_mnk[1]: + # Convert to C type + acc_vec = tiled_copy_r2s.retile(tTR_rAcc).load() + acc_vec = acc_vec.to(self.c_dtype) + tRS_rC.store(acc_vec) + num_prev_subtiles += 1 + c_buffer = num_prev_subtiles % self.num_c_stage + # Store C to shared memory + cute.copy( + tiled_copy_r2s, + tRS_rC, + tRS_sC[(None, None, None, c_buffer)], + ) + # Fence and barrier to make sure shared memory store is visible to TMA store + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) + self.epilog_sync_barrier.arrive_and_wait() + # TMA store C to global memory + if warp_idx == self.epilog_warp_id[0]: + cute.copy( + tma_atom_c, + bSG_sC[(None, c_buffer)], + bSG_gC[(None, subtile_idx)], + ) + c_pipeline.producer_commit() + c_pipeline.producer_acquire() + self.epilog_sync_barrier.arrive_and_wait() + else: + # Convert to C type + acc_vec = tTR_rAcc.load() + acc_vec = acc_vec.to(self.c_dtype) + tTR_rC.store(acc_vec) + # Compute predicate for SIMT store + tCpC = cute.make_rmem_tensor( + cute.make_layout(tTR_rC.shape), + cutlass.Boolean, + ) + m_thr_slice = m_thr_offset[(None, None, None, subtile_idx)] + for i in cutlass.range(cute.size(tCpC), unroll_full=True): + tCpC[i] = ( + m_thr_slice[(i)][0] + + work_tile.cta_coord_m * self.cta_tile_shape_mnk[0] + < tensor_c.shape[0] + ) and (m_thr_slice[(i)][1] < work_tile.distance_to_boundary) + # Store C to global memory + cute.copy( + simt_atom, + cute.flatten(tTR_rC), + cute.flatten(tTR_gC[(None, None, None, subtile_idx)]), + pred=cute.flatten(tCpC), + ) + # Async arrive accumulator buffer empty + with cute.arch.elect_one(): + acc_pipeline.consumer_release(acc_consumer_state) + acc_consumer_state.advance() + # Advance to next tile + tile_info_pipeline.consumer_wait(tile_info_consumer_state) + work_tile = mixed_input_utils.make_contiguous_group_work_tile_info( + group_count, sTile_info[(None, tile_info_consumer_state.index)] + ) + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) + tile_info_pipeline.consumer_release(tile_info_consumer_state) + tile_info_consumer_state.advance() + + # Dealloc the tensor memory buffer + tmem.relinquish_alloc_permit() + self.epilog_sync_barrier.arrive_and_wait() + tmem.free(tmem_ptr) + c_pipeline.producer_tail() + + @staticmethod + def _compute_stages_and_tmem_cols( + tiled_mma: cute.TiledMma, + mma_tiler_mnk: tuple[int, int, int], + cta_tile_shape_mnk: tuple[int, int, int], + epi_tile: cute.Tile, + a_dtype: type[cutlass.Numeric], + b_dtype: type[cutlass.Numeric], + c_dtype: type[cutlass.Numeric], + c_layout: utils.LayoutEnum, + transform_a_source: tcgen05.OperandSource, + scale_granularity_m: int, + scale_granularity_k: int, + smem_buffer_align_bytes: int, + scale_mode: TransformMode, + ) -> tuple[int, int, int, int, int, int, int, int]: + """ + Compute pipeline stages and TMEM column allocation configurations. + + This method calculates the number of pipeline stages for different operations + (tile_info, load2trans, trans2mma, accumulator, etc.) and determines TMEM column allocation + based on available memory resources and tile configuration. + + :param tiled_mma: The tiled MMA object defining the core computation. + :type tiled_mma: cute.TiledMma + :param mma_tiler_mnk: The shape (M, N, K) of the MMA tiler. + :type mma_tiler_mnk: tuple[int, int, int] + :param cta_tile_shape_mnk: The shape (M, N, K) of the CTA tile. + :type cta_tile_shape_mnk: tuple[int, int, int] + :param epi_tile: The epilogue tile shape. + :type epi_tile: cute.Tile + :param a_dtype: Data type of operand A. + :type a_dtype: type[cutlass.Numeric] + :param b_dtype: Data type of operand B. + :type b_dtype: type[cutlass.Numeric] + :param c_dtype: Data type of operand C. + :type c_dtype: type[cutlass.Numeric] + :param c_layout: Layout enum of operand C. + :type c_layout: utils.LayoutEnum + :param transform_a_source: The source of the transformed A tensor. + :type transform_a_source: tcgen05.OperandSource + :param scale_granularity_m: The granularity of the scale tensor along the M mode. + :type scale_granularity_m: int + :param scale_granularity_k: The granularity of the scale tensor along the K mode. + :type scale_granularity_k: int + :param smem_buffer_align_bytes: The alignment of the shared memory buffer. + :type smem_buffer_align_bytes: int + :param scale_mode: The transform mode. + :type scale_mode: TransformMode + + :return: A tuple containing the number of stages for: + (load2trans, scale_load2trans, transform2mma, accumulator, c, tile_info, tmem_acc_cols, tmem_a_cols) + :rtype: tuple[int, int, int, int, int, int, int] + - num_load2trans_stage: Stages for load-to-transform A and B tensors pipeline + - num_scale_load2trans_stage: Stages for scale load-to-transform A tensor pipeline + - num_trans2mma_stage: Stages for transform-to-MMA pipeline + - num_acc_stage: Stages for accumulator-to-epilogue pipeline + - num_c_stage: Stages for epilogue-to-output C pipeline + - num_tile_info_stage: Stages for buffers storing tile info + - num_acc_tmem_cols: TMEM columns for accumulator + - num_a_tmem_cols: TMEM columns for transformed A tensor + """ + # Compute tmem columns required for accumulator + acc_shape = tiled_mma.partition_shape_C(mma_tiler_mnk[:2]) + tCtAcc_stage1 = tiled_mma.make_fragment_C(cute.append(acc_shape, 1)) + num_tmem_acc_col_per_stage = utils.get_num_tmem_alloc_cols(tCtAcc_stage1, True) + # Heuristic to decide the number of stages for accumulator + sm100_tmem_columns = cute.arch.get_max_tmem_alloc_cols("sm_100") + accumulator_stage_count = sm100_tmem_columns // num_tmem_acc_col_per_stage + if transform_a_source == tcgen05.OperandSource.TMEM: + if num_tmem_acc_col_per_stage < 128: + accumulator_stage_count = 3 + elif num_tmem_acc_col_per_stage < 256: + accumulator_stage_count = 2 + else: + accumulator_stage_count = 1 + # transformed A in 16bit, thus 1 tmem column could hold 2 elements + num_elts_per_tmem_col = 32 // tiled_mma.op.a_dtype.width + num_tmem_cols_a_per_stage = cute.round_up( + ( + cta_tile_shape_mnk[2] // num_elts_per_tmem_col + if transform_a_source == tcgen05.OperandSource.TMEM + else 0 + ), + 4, + ) + + bytes_per_pipeline_stage = 16 + # By default, we use 2 stages for tile info + num_tile_info_stage = 2 + tile_info_bytes = ( + cute.size_in_bytes(cute.Int32, cute.make_layout((4, num_tile_info_stage))) + + bytes_per_pipeline_stage * num_tile_info_stage + ) + + c_stage_count = 2 + c_smem_layout_staged_one = sm100_utils.make_smem_layout_epi( + c_dtype, + c_layout, + epi_tile, + 1, + ) + c_bytes_per_stage = cute.size_in_bytes(c_dtype, c_smem_layout_staged_one) + c_bytes = c_bytes_per_stage * c_stage_count + + smem_capacity = utils.get_smem_capacity_in_bytes("sm_100") + if scale_mode == TransformMode.ConvertOnly: + scale_load2trans_stage_count = 0 + a_scale_bytes_per_stage = 0 + else: + # Ensure we have 4 buffers for scale tiles needed for 1 CTA tile + a_scale_k_mode = max(cta_tile_shape_mnk[2] // scale_granularity_k, 1) + a_scale_m_mode = max(cta_tile_shape_mnk[0] // scale_granularity_m, 1) + scale_load2trans_stage_count = 4 + a_scale_bytes_per_stage = cute.round_up( + cute.size_in_bytes( + tiled_mma.op.a_dtype, + cute.make_layout((a_scale_m_mode, a_scale_k_mode)), + ), + smem_buffer_align_bytes, + ) + a_scale_bytes = ( + a_scale_bytes_per_stage + bytes_per_pipeline_stage + ) * scale_load2trans_stage_count + carveout_smem_bytes = ( + bytes_per_pipeline_stage * accumulator_stage_count + + a_scale_bytes + + c_bytes + + tile_info_bytes + ) + + # Compute transform stages if A is in TMEM + num_tmem_acc_cols = cute.round_up( + accumulator_stage_count * num_tmem_acc_col_per_stage, 4 + ) + + transform2mma_stage_count_a_source_tmem_potential = ( + (sm100_tmem_columns - num_tmem_acc_cols) // num_tmem_cols_a_per_stage + if transform_a_source == tcgen05.OperandSource.TMEM + else -1 + ) + if ( + transform_a_source == tcgen05.OperandSource.TMEM + and transform2mma_stage_count_a_source_tmem_potential <= 0 + ): + raise ValueError("Not enough TMEM capacity for selected tile size") + a_load_bytes_per_stage = cute.round_up( + cute.size_in_bytes( + a_dtype, + cute.make_layout((cta_tile_shape_mnk[0], cta_tile_shape_mnk[2])), + ), + smem_buffer_align_bytes, + ) + b_load_bytes_per_stage = cute.round_up( + cute.size_in_bytes( + b_dtype, + cute.make_layout( + ( + cta_tile_shape_mnk[1] // cute.size(tiled_mma.thr_id), + cta_tile_shape_mnk[2], + ) + ), + ), + smem_buffer_align_bytes, + ) + ab_load_bytes_per_stage = ( + a_load_bytes_per_stage + + b_load_bytes_per_stage + + 2 * bytes_per_pipeline_stage + ) + a_transform_bytes_per_stage = ( + cute.round_up( + cute.size_in_bytes( + tiled_mma.op.a_dtype, + cute.make_layout((cta_tile_shape_mnk[0], cta_tile_shape_mnk[2])), + ), + smem_buffer_align_bytes, + ) + if transform_a_source == tcgen05.OperandSource.SMEM + else 0 + ) + + a_transform_bytes_per_stage = ( + a_transform_bytes_per_stage + bytes_per_pipeline_stage + ) + transform2mma_stage_count_a_source_smem_potential = ( + smem_capacity - carveout_smem_bytes + ) // (ab_load_bytes_per_stage + a_transform_bytes_per_stage) + transform2mma_stage_count = ( + min( + transform2mma_stage_count_a_source_tmem_potential, + transform2mma_stage_count_a_source_smem_potential, + ) + if transform_a_source == tcgen05.OperandSource.TMEM + else transform2mma_stage_count_a_source_smem_potential + ) + load2transform_stage_count = ( + smem_capacity + - carveout_smem_bytes + - (transform2mma_stage_count * a_transform_bytes_per_stage) + ) // ab_load_bytes_per_stage + if ( + load2transform_stage_count < 2 + or transform2mma_stage_count < 2 + or accumulator_stage_count < 1 + ): + raise ValueError("Not enough SMEM or TMEM capacity for selected tile size") + num_tmem_a_cols = transform2mma_stage_count * num_tmem_cols_a_per_stage + # Check if we can increase c_stage_count with leftover smem + c_stage_count += ( + smem_capacity + - load2transform_stage_count * ab_load_bytes_per_stage + - transform2mma_stage_count * a_transform_bytes_per_stage + - scale_load2trans_stage_count * a_scale_bytes_per_stage + - c_bytes + ) // c_bytes_per_stage + + return ( + load2transform_stage_count, + scale_load2trans_stage_count, + transform2mma_stage_count, + accumulator_stage_count, + c_stage_count, + num_tile_info_stage, + num_tmem_acc_cols, + num_tmem_a_cols, + ) + + @staticmethod + def _compute_grid( + c: cute.Tensor, + cta_tile_shape_mnk: tuple[int, int, int], + cluster_shape_mn: tuple[int, int], + max_active_clusters: cutlass.Constexpr, + ) -> tuple[utils.PersistentTileSchedulerParams, tuple[int, int, int]]: + """ + Use persistent tile scheduler to compute the grid size for the output tensor C. + """ + c_shape = cute.slice_(cta_tile_shape_mnk, (None, None, 0)) + gc = cute.zipped_divide(c, tiler=c_shape) + num_ctas_mnl = gc[(0, (None, None, None))].shape + cluster_shape_mnl = (*cluster_shape_mn, 1) + + tile_sched_params = utils.PersistentTileSchedulerParams( + num_ctas_mnl, cluster_shape_mnl + ) + grid = (cluster_shape_mn[0], cluster_shape_mn[1], max_active_clusters) + + return tile_sched_params, grid + + def can_implement( + mnkl: tuple[int, int, int, int], + a_dtype: type[cutlass.Numeric], + b_dtype: type[cutlass.Numeric], + c_dtype: type[cutlass.Numeric], + a_major: str, + b_major: str, + c_major: str, + scale_granularity_m: int, + scale_granularity_k: int, + mma_tiler: tuple[int, int, int], + cluster_shape_mn: tuple[int, int], + use_2cta_instrs: bool, + ) -> bool: + """ + Check if the kernel can be implemented for the given tensor shapes and data types. + """ + m, n, k, l = mnkl + + if not mixed_input_utils.is_valid_mma_tiler_and_cluster_shape( + mma_tiler, cluster_shape_mn, use_2cta_instrs + ): + return False + if not mixed_input_utils.is_valid_scale_granularity( + scale_granularity_m, scale_granularity_k, a_dtype, k, mma_tiler[2] + ): + return False + if not mixed_input_utils.is_valid_tensor_alignment( + m, + n, + k, + a_dtype, + b_dtype, + c_dtype, + b_dtype, + a_major, + b_major, + c_major, + mma_tiler, + use_2cta_instrs, + cluster_shape_mn, + scale_granularity_m, + scale_granularity_k, + ): + return False + return True + + +def get_advanced_compiler_control_path(): + """ + Return the path to the advanced compiler control file of this example. If not found, return None. + """ + import os + + need_advanced_compiler_control = False + try: + from cutlass import CUDA_VERSION + + if CUDA_VERSION.major == 13 and CUDA_VERSION.minor == 1: + need_advanced_compiler_control = True + except ImportError: + pass + + if not need_advanced_compiler_control: + return None + # Get the path to the advanced compiler control file + current_dir = os.path.dirname(os.path.abspath(__file__)) + target_path = os.path.join(current_dir, "../../advanced_compiler_control/gemm0.bin") + if os.path.exists(target_path): + print(f"Found advanced compiler control file at {target_path}") + return target_path + else: + return None + + +def run( + mnkl: tuple[int, int, int, int], + scale_granularity_m: int, + scale_granularity_k: int, + a_dtype: type[cutlass.Numeric], + b_dtype: type[cutlass.Numeric], + c_dtype: type[cutlass.Numeric], + acc_dtype: type[cutlass.Numeric], + a_major: str, + b_major: str, + c_major: str, + mma_tiler_mnk: tuple[int, int, int], + cluster_shape_mn: tuple[int, int], + use_2cta_instrs: bool, + tolerance: float, + warmup_iterations: int = 0, + iterations: int = 1, + skip_ref_check: bool = False, + uniform_group_sizes: bool = False, + use_cold_l2: bool = False, + **kwargs, +) -> None: + """ + Run the mixed-input GEMM kernel with specified parameters. + + This function creates tensors, validates parameters, executes the kernel, + optionally compares results with a reference implementation and reports + kernel execution time. + """ + m, n, k, l = mnkl + import torch + + if not torch.cuda.is_available(): + raise ValueError("CUDA is not available") + + # Check if given configuration is supported + if not GroupedMixedInputGemmKernel.can_implement( + mnkl, + a_dtype, + b_dtype, + c_dtype, + a_major, + b_major, + c_major, + scale_granularity_m, + scale_granularity_k, + mma_tiler_mnk, + cluster_shape_mn, + use_2cta_instrs, + ): + raise ValueError("GEMM configuration not supported") + + # Get current CUDA stream from PyTorch + torch_stream = torch.cuda.current_stream() + # Get the raw stream pointer as a CUstream + current_stream = cuda.CUstream(torch_stream.cuda_stream) + + group_count = l + shuffle_a = mixed_input_utils.is_shuffle_a( + a_major, k, a_dtype, b_dtype, scale_granularity_k + ) + # shuffle is supported since CUDA 13.1 + shuffle_supported = False + try: + from cutlass import CUDA_VERSION + + if CUDA_VERSION.major > 13 or ( + CUDA_VERSION.major == 13 and CUDA_VERSION.minor >= 1 + ): + shuffle_supported = True + except ImportError: + pass + + shuffle_a = shuffle_a and shuffle_supported + mixed_input_gemm = GroupedMixedInputGemmKernel( + scale_granularity_m, + scale_granularity_k, + acc_dtype, + use_2cta_instrs, + mma_tiler_mnk, + cluster_shape_mn, + group_count, + shuffle_a, + ) + torch.manual_seed(2025) + ( + a_tensor, + a_scale_tensor, + b_tensor, + cumsum_tensor, + c_tensor, + a_torch_cpu, + a_scale_torch_cpu, + b_torch_cpu, + cumsum_torch_cpu, + c_torch_gpu, + ) = create_tensors( + l, + m, + n, + k, + a_major, + b_major, + c_major, + a_dtype, + b_dtype, + c_dtype, + shuffle_a, + scale_granularity_m, + scale_granularity_k, + uniform_group_sizes, + ) + + max_active_clusters = utils.HardwareInfo().get_max_active_clusters( + cluster_shape_mn[0] * cluster_shape_mn[1], + ) + advanced_compiler_options = None + advanced_compiler_control_path = get_advanced_compiler_control_path() + if advanced_compiler_control_path: + advanced_compiler_options = ( + f"--ptxas-options '--apply-controls={advanced_compiler_control_path}'" + ) + + compiled_kernel = cute.compile( + mixed_input_gemm, + a_tensor, + a_scale_tensor, + b_tensor, + cumsum_tensor, + c_tensor, + max_active_clusters, + current_stream, + options=advanced_compiler_options, + ) + + if not skip_ref_check: + compiled_kernel( + a_tensor, + a_scale_tensor, + b_tensor, + cumsum_tensor, + c_tensor, + current_stream, + ) + run_ref_and_compare( + a_torch_cpu, + b_torch_cpu, + a_scale_torch_cpu, + cumsum_torch_cpu, + c_torch_gpu, + c_dtype, + tolerance, + ) + + # Early return if no performance measurement is needed + if iterations <= 0: + return + + def generate_tensors(): + ( + a_tensor, + a_scale_tensor, + b_tensor, + cumsum_tensor, + c_tensor, + a_torch_cpu, + a_scale_torch_cpu, + b_torch_cpu, + cumsum_torch_cpu, + c_torch_gpu, + ) = create_tensors( + l, + m, + n, + k, + a_major, + b_major, + c_major, + a_dtype, + b_dtype, + c_dtype, + shuffle_a, + scale_granularity_m, + scale_granularity_k, + uniform_group_sizes, + ) + return testing.JitArguments( + a_tensor, a_scale_tensor, b_tensor, cumsum_tensor, c_tensor, current_stream + ) + + workspace_count = 1 + if use_cold_l2: + one_workspace_bytes = ( + a_torch_cpu.numel() * a_torch_cpu.element_size() + + b_torch_cpu.numel() * b_torch_cpu.element_size() + + c_torch_gpu.numel() * c_torch_gpu.element_size() + + a_scale_torch_cpu.numel() * a_scale_torch_cpu.element_size() + if a_scale_torch_cpu is not None + else 0 + ) + workspace_count = testing.get_workspace_count( + one_workspace_bytes, warmup_iterations, iterations + ) + + exec_time = testing.benchmark( + compiled_kernel, + workspace_generator=generate_tensors, + workspace_count=workspace_count, + stream=current_stream, + warmup_iterations=warmup_iterations, + iterations=iterations, + ) + + return exec_time # Return execution time in microseconds + + +if __name__ == "__main__": + + def parse_comma_separated_ints(s: str) -> tuple[int, ...]: + try: + return tuple(int(x.strip()) for x in s.split(",")) + except ValueError: + raise argparse.ArgumentTypeError( + "Invalid format. Expected comma-separated integers." + ) + + parser = argparse.ArgumentParser() + parser.add_argument( + "--mnkl", type=parse_comma_separated_ints, default=(128, 128, 128, 1) + ) + parser.add_argument( + "--mma_tiler_mnk", type=parse_comma_separated_ints, default=(128, 128, 128) + ) + parser.add_argument( + "--cluster_shape_mn", type=parse_comma_separated_ints, default=(1, 1) + ) + parser.add_argument( + "--use_2cta_instrs", + action="store_true", + help="Enable 2CTA MMA instructions feature", + ) + parser.add_argument( + "--a_dtype", + type=cutlass.dtype, + default=cutlass.Int4, + choices=[cutlass.Int8, cutlass.Uint8, cutlass.Int4], + ) + parser.add_argument( + "--b_dtype", + type=cutlass.dtype, + default=cutlass.BFloat16, + choices=[cutlass.BFloat16, cutlass.Float16], + ) + parser.add_argument("--c_dtype", type=cutlass.dtype, default=cutlass.BFloat16) + parser.add_argument("--acc_dtype", type=cutlass.dtype, default=cutlass.Float32) + parser.add_argument("--a_major", choices=["k", "m"], type=str, default="m") + parser.add_argument("--b_major", choices=["k", "n"], type=str, default="k") + parser.add_argument("--c_major", choices=["n", "m"], type=str, default="n") + parser.add_argument( + "--scale_granularity_m", + type=int, + default=1, + help="Scale granularity along M dimension.", + ) + parser.add_argument( + "--scale_granularity_k", + type=int, + default=128, + help="Scale granularity along K dimension.", + ) + parser.add_argument( + "--tolerance", type=float, default=1e-01, help="Tolerance for validation" + ) + parser.add_argument( + "--warmup_iterations", type=int, default=0, help="Warmup iterations" + ) + parser.add_argument( + "--iterations", + type=int, + default=1, + help="Number of iterations to run the kernel", + ) + parser.add_argument( + "--skip_ref_check", action="store_true", help="Skip reference checking" + ) + parser.add_argument( + "--uniform_group_sizes", action="store_true", help="Use uniform group sizes" + ) + args = parser.parse_args() + run( + args.mnkl, + args.scale_granularity_m, + args.scale_granularity_k, + args.a_dtype, + args.b_dtype, + args.c_dtype, + args.acc_dtype, + args.a_major, + args.b_major, + args.c_major, + args.mma_tiler_mnk, + args.cluster_shape_mn, + args.use_2cta_instrs, + args.tolerance, + args.warmup_iterations, + args.iterations, + args.skip_ref_check, + args.uniform_group_sizes, + ) + print("PASS") diff --git a/examples/python/CuTeDSL/blackwell/mixed_input_gemm/grouped_mixed_input_gemm_acc_scale.py b/examples/python/CuTeDSL/blackwell/mixed_input_gemm/grouped_mixed_input_gemm_acc_scale.py new file mode 100644 index 00000000..472816b6 --- /dev/null +++ b/examples/python/CuTeDSL/blackwell/mixed_input_gemm/grouped_mixed_input_gemm_acc_scale.py @@ -0,0 +1,2502 @@ +# 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. + +import argparse +from math import log2, ceil +from typing import Union +import os +import sys + +import torch +import cuda.bindings.driver as cuda + +import cutlass +import cutlass.cute as cute +import cutlass.pipeline as pipeline +from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait +import cutlass.utils as utils +import cutlass.utils.blackwell_helpers as sm100_utils +import cutlass.utils.mixed_input_helpers as mixed_input_utils +from cutlass.utils.mixed_input_helpers import TransformMode +import cutlass.cute.testing as testing +from cutlass.cute.nvgpu import cpasync, tcgen05 + +if __name__ == "__main__": + current_dir = os.path.dirname(os.path.abspath(__file__)) + sys.path.insert(0, os.path.join(current_dir, "../..")) + +from blackwell.mixed_input_gemm.mixed_input_host_utils import ( + create_tensors_for_contiguous_grouped_mixed_input_gemm as create_tensors, + run_contiguous_grouped_ref_and_compare as run_ref_and_compare, +) + +""" +A mixed-input grouped GEMM example for the NVIDIA Blackwell SM100 architecture using CUTE DSL. + +Compared to the grouped_mixed_input_gemm.py example, this acc_scale example demonstrates a different implementation by +performing the scaling step on the accumulator instead of the input A tensor. The original computation is +``` +C = (type_convert(A) * scale) x B. +``` +In this example, we swap the scaling step and MMA and reformulate the computation as follows: +``` +C = scale * (type_convert(A) x B). +``` +The reformulation is valid when `scale_granularity_k` is an exact multiple of `mma_tiler_mnk[2]`. With this change, +the number of operations in the scaling step changes from `CTA_TILE_M * CTA_TILE_K` to `CTA_TILE_M * CTA_TILE_N`. +When `CTA_TILE_N` is smaller than `CTA_TILE_K` (common in decoding), the acc_scale implementation can +be more efficient due to fewer scaling operations. +Only convert-scale mode is supported in this example; convert-only mode has no scaling. +Other than the above changes, the computation flow is the same as the grouped_mixed_input_gemm.py example. + +To run this example: + +.. code-block:: bash + + python examples/blackwell/mixed_input_gemm/grouped_mixed_input_gemm_acc_scale.py \ + --a_dtype Int4 --b_dtype BFloat16 \ + --scale_granularity_m 1 --scale_granularity_k 256 \ + --c_dtype BFloat16 --acc_dtype Float32 \ + --mma_tiler_mnk 256,16,128 --cluster_shape_mn 2,1 \ + --use_2cta_instrs --mnkl 1024,8192,6144,16 \ + +To collect performance with NCU profiler: + +.. code-block:: bash + + ncu python examples/blackwell/mixed_input_gemm/grouped_mixed_input_gemm_acc_scale.py \ + --a_dtype Int4 --b_dtype BFloat16 \ + --scale_granularity_m 1 --scale_granularity_k 256 \ + --c_dtype BFloat16 --acc_dtype Float32 \ + --mma_tiler_mnk 128,8,256 --cluster_shape_mn 1,1 \ + --mnkl 4096,8,8192,32 \ + --warmup_iterations 1 --iterations 10 --skip_ref_check +""" + + +class GroupedMixedInputGemmAccScaleKernel: + """ + Mixed-input grouped GEMM kernel with scaling applied on accumulator for NVIDIA Blackwell SM100 architecture. + + This kernel supports GEMM operations where input tensors A and B have different data types, with tensor A + being converted to the precision of tensor B before matrix multiplication. The scaling is applied on + the accumulator instead of the input A tensor to reduce the number of operations in the scaling step. + Tensor A is in shape of [M, K, L] with L being the number of groups. Tensor B is in shape of [N, K] and a group search algorithm + is applied along the N mode to find the group index for each CTA tile. A cumsum tensor provides the offsets of each group along the N mode. + + :param scale_granularity_m: Number of elements sharing the same scale factor along the M mode + :type scale_granularity_m: int + :param scale_granularity_k: Number of elements sharing the same scale factor along the K mode + :type scale_granularity_k: int + :param acc_dtype: Data type for accumulation during computation + :type acc_dtype: type[cutlass.Numeric] + :param use_2cta_instrs: Whether to use CTA group 2 for advanced thread cooperation + :type use_2cta_instrs: bool + :param mma_tiler_mnk: Shape of the Matrix Multiply-Accumulate (MMA) tile (M, N, K) + :type mma_tiler_mnk: tuple[int, int, int] + :param cluster_shape_mn: Cluster dimensions (M,N) for parallel processing + :type cluster_shape_mn: tuple[int, int] + :param group_count: The total number of groups + :type group_count: int + :param shuffle_a: Whether to use shuffle intrinsic for int4-to-bf16 conversion + :type shuffle_a: bool + """ + + def __init__( + self, + scale_granularity_m: int, + scale_granularity_k: int, + acc_dtype: type[cutlass.Numeric], + use_2cta_instrs: bool, + mma_tiler_mnk: tuple[int, int, int], + cluster_shape_mn: tuple[int, int], + group_count: int, + shuffle_a: bool, + ): + """ + Initializes the mixed-input GEMM kernel with a specified configuration. + """ + # Scale granularity defines how many elements share the same scale factor + # along the M and K modes. + self.scale_granularity_m = scale_granularity_m + self.scale_granularity_k = scale_granularity_k + # Set transform mode + if cutlass.const_expr( + self.scale_granularity_m == 0 and self.scale_granularity_k == 0 + ): + # Acc-update kernel variant is only for convert-scale mode. + raise ValueError("convert-only mode is not supported for acc-scale kernel") + self.scale_mode = TransformMode.ConvertScale + # scale_granularity_k must be exactly multiple of CTA tile shape K to allow acc-update recipe. + if cutlass.const_expr(self.scale_granularity_k % mma_tiler_mnk[2] != 0): + raise ValueError( + "scale_granularity_k must be exactly multiple of CTA tile shape K" + ) + + self.group_count = group_count + self.acc_dtype = acc_dtype + self.use_2cta_instrs = use_2cta_instrs + self.cluster_shape_mn = cluster_shape_mn + self.mma_tiler = mma_tiler_mnk + self.shuffle_a = shuffle_a + self.cta_group = ( + tcgen05.CtaGroup.TWO if self.use_2cta_instrs else tcgen05.CtaGroup.ONE + ) + # transformation ktile loop unrolling factor + self.transform_k_tile_unroll_factor = 2 + # Set specialized warp ids + self.epilog_warp_id = ( + 0, + 1, + 2, + 3, + ) + self.mma_warp_id = 4 + self.tma_warp_id = 5 + self.scale_tma_warp_id = 6 + # Schedule warp to do the group search + self.schedule_warp_id = 7 + self.transform_warp_id = ( + 8, + 9, + 10, + 11, + ) + # Define expected register count for different warps + # Reserve more registers for transformation and epilogue warps + self.num_regs_epilogue_warps = 168 + self.num_regs_mma_warp = 80 + self.num_regs_tma_warps = 72 + self.num_regs_transform_warps = 240 + self.num_regs_schedule_warp = 64 + self.threads_per_cta = 32 * ( + max( + ( + self.mma_warp_id, + self.tma_warp_id, + self.scale_tma_warp_id, + *self.epilog_warp_id, + *self.transform_warp_id, + ) + ) + + 1 + ) + + # Set barrier id for cta sync, epilogue sync, tmem ptr sync, and transform sync + self.epilog_sync_barrier = pipeline.NamedBarrier( + 1, 32 * len(self.epilog_warp_id) + ) + self.tmem_ptr_sync_barrier = pipeline.NamedBarrier(2, self.threads_per_cta) + self.cta_sync_barrier = pipeline.NamedBarrier(3, self.threads_per_cta) + self.sched_sync_barrier = pipeline.NamedBarrier(4, 32) + + self.smem_buffer_align_bytes = 1024 + + def _setup_attributes(self): + """ + Set up configurations that are dependent on GEMM inputs + """ + # Deduce where the transformed A tensor is stored, shared memory(SMEM) or tensor memory(TMEM) + self.transform_a_source = mixed_input_utils.get_transform_a_source( + self.a_major_mode + ) + tiled_mma = sm100_utils.make_trivial_tiled_mma( + self.mma_dtype, + self.a_major_mode, + self.b_major_mode, + self.acc_dtype, + self.cta_group, + self.mma_tiler[:2], + self.transform_a_source, + ) + self.cta_tile_shape_mnk = ( + self.mma_tiler[0] // cute.size(tiled_mma.thr_id.shape), + self.mma_tiler[1], + self.mma_tiler[2], + ) + self.cluster_tile_shape_mnk = ( + self.cluster_shape_mn[0] * self.cta_tile_shape_mnk[0], + self.cluster_shape_mn[1] * self.cta_tile_shape_mnk[1], + self.cta_tile_shape_mnk[2], + ) + self.cluster_layout_vmnk = cute.tiled_divide( + cute.make_layout((*self.cluster_shape_mn, 1)), + (tiled_mma.thr_id.shape,), + ) + self.num_mcast_ctas_a = cute.size(self.cluster_layout_vmnk.shape[2]) + self.num_mcast_ctas_b = cute.size(self.cluster_layout_vmnk.shape[1]) + self.is_a_mcast = self.num_mcast_ctas_a > 1 + self.is_b_mcast = self.num_mcast_ctas_b > 1 + + self.epi_tile = sm100_utils.compute_epilogue_tile_shape( + self.cta_tile_shape_mnk, + self.use_2cta_instrs, + self.c_layout, + self.c_dtype, + ) + + # Compute tensor memory(TMEM) columns and stages for each pipeline + ( + self.num_load2trans_stage, + self.num_scale_load2accu_stage, + self.num_trans2mma_stage, + self.num_acc_stage, + self.num_c_stage, + self.num_tile_info_stage, + self.num_acc_tmem_cols, + self.num_a_tmem_cols, + ) = self._compute_stages_and_tmem_cols( + tiled_mma, + self.mma_tiler, + self.cta_tile_shape_mnk, + self.epi_tile, + self.a_dtype, + self.b_dtype, + self.c_dtype, + self.c_layout, + self.transform_a_source, + self.scale_granularity_m, + self.scale_granularity_k, + self.smem_buffer_align_bytes, + ) + + # Align TMEM columns for allocation + # TMEM allocation requires power-of-2 column alignment + # and must meet minimum allocation requirements + self.num_tmem_alloc_cols = cute.round_up( + self.num_acc_tmem_cols + self.num_a_tmem_cols, + cute.arch.get_min_tmem_alloc_cols("sm_100"), + ) + self.num_tmem_alloc_cols = 2 ** (ceil(log2(self.num_tmem_alloc_cols))) + # Get smem layout for C tensor + self.c_smem_layout_staged = sm100_utils.make_smem_layout_epi( + self.c_dtype, + self.c_layout, + self.epi_tile, + self.num_c_stage, + ) + # Get smem layout for A, transformed A, and B + ( + self.smem_layout_a, + self.smem_layout_a_transform, + self.smem_layout_b, + ) = mixed_input_utils.compute_smem_layout( + tiled_mma, + self.mma_tiler, + self.a_dtype, + self.b_dtype, + self.num_load2trans_stage, + self.num_trans2mma_stage, + ) + # Check if stages match the requirements for unrolling + if ( + self.num_scale_load2accu_stage < self.transform_k_tile_unroll_factor + or self.num_trans2mma_stage < self.transform_k_tile_unroll_factor + ): + raise ValueError("Not enough SMEM capacity for selected tile size") + # Get scale tile shape and smem layout for scale tensor + # ((M_SHARING_SCALE, NUM_SCALES_M),(K_SHARING_SCALE, NUM_SCALES_K), STAGES) + ( + self.scale_tile_shape, + self.smem_layout_scale_per_stage, + self.smem_layout_scale, + ) = mixed_input_utils.get_smem_layout_scale( + self.mma_tiler, + self.use_2cta_instrs, + self.scale_granularity_m, + self.scale_granularity_k, + self.scale_major_mode, + self.a_scale_dtype, + self.num_scale_load2accu_stage, + ) + + def _validate_inputs( + self, + a: cute.Tensor, + a_scale: cute.Tensor, + b: cute.Tensor, + c: cute.Tensor, + ) -> None: + """ + Validates input tensors and their properties. + """ + # Validate scale tensor major mode + if cutlass.const_expr( + utils.LayoutEnum.from_tensor(a_scale).mma_major_mode() + != tcgen05.OperandMajorMode.MN + ): + raise ValueError("scale_major_mode must be M-major") + + @cute.jit + def __call__( + self, + a: cute.Tensor, + a_scale: cute.Tensor, + b: cute.Tensor, + cumsum: cute.Tensor, + c: cute.Tensor, + max_active_clusters: cutlass.Constexpr, + stream: cuda.CUstream, + ): + """ + Executes the Mixed Input Grouped GEMM operation. + """ + self.a_dtype: type[cutlass.Numeric] = a.element_type + self.a_scale_dtype: type[cutlass.Numeric] = a_scale.element_type + self.b_dtype: type[cutlass.Numeric] = b.element_type + self.c_dtype: type[cutlass.Numeric] = c.element_type + self.mma_dtype = self.b_dtype + + self.a_major_mode = utils.LayoutEnum.from_tensor(a).mma_major_mode() + self.scale_major_mode = utils.LayoutEnum.from_tensor(a_scale).mma_major_mode() + self.b_major_mode = utils.LayoutEnum.from_tensor(b).mma_major_mode() + self.c_layout = utils.LayoutEnum.from_tensor(c) + # Get gmem layout for scale tensor + self.gmem_layout_scale = mixed_input_utils.get_gmem_layout_scale( + a.shape, + self.scale_granularity_m, + self.scale_granularity_k, + self.scale_major_mode, + ) + + # Validate inputs + self._validate_inputs(a, a_scale, b, c) + + # Setup attributes that dependent on gemm inputs + self._setup_attributes() + + tiled_mma = sm100_utils.make_trivial_tiled_mma( + self.mma_dtype, + self.a_major_mode, + self.b_major_mode, + self.acc_dtype, + self.cta_group, + self.mma_tiler[:2], + self.transform_a_source, + ) + # Set up gmem copy atoms for A, scale, and B + a_op = mixed_input_utils.get_tma_atom_kind( + self.is_a_mcast, self.use_2cta_instrs, is_b=False + ) + b_op = mixed_input_utils.get_tma_atom_kind( + self.is_b_mcast, self.use_2cta_instrs, is_b=True + ) + a_scale_op = a_op + # Deduce TMA copy atom and TMA tensor for A, scale, and B + smem_layout_a_per_stage = cute.slice_(self.smem_layout_a, (None, None, None, 0)) + tma_atom_a, tma_tensor_a = cute.nvgpu.make_tiled_tma_atom_A( + a_op, + a, + smem_layout_a_per_stage, + self.mma_tiler, + tiled_mma, + self.cluster_layout_vmnk.shape, + internal_type=( + cutlass.TFloat32 if a.element_type is cutlass.Float32 else None + ), + ) + + # Partition smem layout for scale tensor to make it compatible with TMA atom + smem_layout_for_tma_atom = cute.get( + tiled_mma._thrfrg_A(self.smem_layout_scale_per_stage.outer), mode=[1] + ) + # ((MMA_M, MMA_K), REST_M, REST_K) + smem_layout_for_tma_atom = cute.dice( + smem_layout_for_tma_atom, + (1, (1,) * cute.rank(self.smem_layout_scale_per_stage.outer)), + ) + tma_atom_scale, tma_tensor_scale = cute.nvgpu.make_tiled_tma_atom_A( + a_scale_op, + cute.make_tensor(a_scale.iterator, self.gmem_layout_scale), + smem_layout_for_tma_atom, + # (SCALE_M, 1, SCALE_K) + (self.scale_tile_shape[0], 1, self.scale_tile_shape[1]), + tiled_mma, + self.cluster_layout_vmnk.shape, + internal_type=( + cutlass.TFloat32 if a_scale.element_type is cutlass.Float32 else None + ), + ) + + smem_layout_b_per_stage = cute.slice_(self.smem_layout_b, (None, None, None, 0)) + tma_atom_b, tma_tensor_b = cute.nvgpu.make_tiled_tma_atom_B( + b_op, + b, + smem_layout_b_per_stage, + self.mma_tiler, + tiled_mma, + self.cluster_layout_vmnk.shape, + internal_type=( + cutlass.TFloat32 if b.element_type is cutlass.Float32 else None + ), + ) + + # Calculate copy size for tensor A, B, and scale + a_copy_size = cute.size_in_bytes(self.a_dtype, smem_layout_a_per_stage) + b_copy_size = cute.size_in_bytes(self.b_dtype, smem_layout_b_per_stage) + a_scale_copy_size = cute.size_in_bytes( + self.a_scale_dtype, self.smem_layout_scale_per_stage + ) + + self.num_tma_load_bytes_a = a_copy_size + self.num_tma_load_bytes_b = b_copy_size * cute.size(tiled_mma.thr_id.shape) + self.num_tma_load_bytes_scale = a_scale_copy_size + self.tile_sched_params, grid = self._compute_grid( + c, + self.cta_tile_shape_mnk, + self.cluster_shape_mn, + max_active_clusters, + ) + + epi_smem_layout = cute.slice_(self.c_smem_layout_staged, (None, None, 0)) + tma_atom_c, tma_tensor_c = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileS2GOp(), + c, + epi_smem_layout, + self.epi_tile, + ) + + # Shared memory structure + @cute.struct + class SharedStorage: + # buffer holding group search results + tile_info: cute.struct.MemRange[cutlass.Int32, 4 * self.num_tile_info_stage] + a_load2trans_full_mbar_ptr: cute.struct.MemRange[ + cutlass.Int64, self.num_load2trans_stage + ] + a_load2trans_empty_mbar_ptr: cute.struct.MemRange[ + cutlass.Int64, self.num_load2trans_stage + ] + a_scale_load2accu_full_mbar_ptr: cute.struct.MemRange[ + cutlass.Int64, self.num_scale_load2accu_stage + ] + a_scale_load2accu_empty_mbar_ptr: cute.struct.MemRange[ + cutlass.Int64, self.num_scale_load2accu_stage + ] + a_trans2mma_full_mbar_ptr: cute.struct.MemRange[ + cutlass.Int64, self.num_trans2mma_stage + ] + a_trans2mma_empty_mbar_ptr: cute.struct.MemRange[ + cutlass.Int64, self.num_trans2mma_stage + ] + b_load2mma_full_mbar_ptr: cute.struct.MemRange[ + cutlass.Int64, self.num_load2trans_stage + ] + b_load2mma_empty_mbar_ptr: cute.struct.MemRange[ + cutlass.Int64, self.num_load2trans_stage + ] + acc_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage] + acc_empty_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage] + tile_info_full_mbar_ptr: cute.struct.MemRange[ + cutlass.Int64, self.num_tile_info_stage + ] + tile_info_empty_mbar_ptr: cute.struct.MemRange[ + cutlass.Int64, self.num_tile_info_stage + ] + tmem_dealloc_mbar_ptr: cutlass.Int64 + tmem_holding_buf: cutlass.Int32 + + self.shared_storage = SharedStorage + + # Launch kernel + self.kernel( + tiled_mma, + tma_atom_a, + tma_tensor_a, + tma_atom_scale, + tma_tensor_scale, + tma_atom_b, + tma_tensor_b, + tma_atom_c, + tma_tensor_c, + c, + cumsum, + self.group_count, + self.cluster_layout_vmnk, + self.smem_layout_a, + self.smem_layout_scale, + self.smem_layout_a_transform, + self.smem_layout_b, + self.c_smem_layout_staged, + self.epi_tile, + self.tile_sched_params, + ).launch( + grid=grid, + block=[self.threads_per_cta, 1, 1], + cluster=(*self.cluster_shape_mn, 1), + min_blocks_per_mp=1, + stream=stream, + ) + return + + # GPU device kernel + @cute.kernel + def kernel( + self, + tiled_mma: cute.TiledMma, + tma_atom_a: cute.CopyAtom, + mA_mkl: cute.Tensor, + tma_atom_s: cute.CopyAtom, + mS_mkl: cute.Tensor, + tma_atom_b: cute.CopyAtom, + mB_nkl: cute.Tensor, + tma_atom_c: cute.CopyAtom, + mC_mnl: cute.Tensor, + tensor_c: cute.Tensor, + cumsum: cute.Tensor, + group_count: cutlass.Constexpr[int], + cluster_layout_vmnk: cute.Layout, + a_smem_layout: cute.ComposedLayout, + scale_smem_layout: cute.ComposedLayout, + a_smem_layout_transform: cute.ComposedLayout, + b_smem_layout: cute.ComposedLayout, + c_smem_layout_staged: cute.ComposedLayout, + epi_tile: cute.Tile, + tile_sched_params: utils.PersistentTileSchedulerParams, + ): + """ + GPU device kernel performing the Persistent Mixed-Input Grouped GEMM computation. + """ + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + tidx, _, _ = cute.arch.thread_idx() + bidx, bidy, bidz = cute.arch.block_idx() + # Prefetch TMA descriptors + if warp_idx == self.epilog_warp_id[0]: + cpasync.prefetch_descriptor(tma_atom_a) + cpasync.prefetch_descriptor(tma_atom_b) + cpasync.prefetch_descriptor(tma_atom_s) + cpasync.prefetch_descriptor(tma_atom_c) + + use_2cta_instrs = cute.size(tiled_mma.thr_id.shape) == 2 + bidx, bidy, bidz = cute.arch.block_idx() + # Compute how many k_tiles share the same scale + num_k_tiles_per_scale = self.scale_granularity_k // self.cta_tile_shape_mnk[2] + + mma_tile_coord_v = bidx % cute.size(tiled_mma.thr_id.shape) + is_leader_cta = mma_tile_coord_v == 0 + cta_rank_in_cluster = cute.arch.make_warp_uniform( + cute.arch.block_idx_in_cluster() + ) + block_in_cluster_coord_vmnk = cluster_layout_vmnk.get_flat_coord( + cta_rank_in_cluster + ) + tidx, _, _ = cute.arch.thread_idx() + + smem = utils.SmemAllocator() + storage = smem.allocate(self.shared_storage) + + # Initialize load2transform pipeline, which tracks the dependencies between TMA's loading + # of A and B, and the transformation of A and MMA's consumption + transform_thread_idx = ( + tidx - 32 * self.transform_warp_id[0] + if tidx >= 32 * self.transform_warp_id[0] + else tidx + ) + a_load2trans_pipeline = pipeline.PipelineTmaAsync.create( + barrier_storage=storage.a_load2trans_full_mbar_ptr.data_ptr(), + num_stages=self.num_load2trans_stage, + producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), + consumer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, + self.num_mcast_ctas_a * len(self.transform_warp_id), + ), + tx_count=self.num_tma_load_bytes_a, + cta_layout_vmnk=cluster_layout_vmnk, + tidx=transform_thread_idx, + mcast_mode_mn=(1, 0), # multicast for A will only happen on the M-mode + defer_sync=True, + ) + # Initialize scale_load2accu pipeline, which tracks the dependencies between TMA's loading + # of scale, and the accumulator update + num_producers_a_scale = self.num_mcast_ctas_a + scale_load2accu_pipeline = pipeline.PipelineTmaAsync.create( + barrier_storage=storage.a_scale_load2accu_full_mbar_ptr.data_ptr(), + num_stages=self.num_scale_load2accu_stage, + producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), + consumer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, + num_producers_a_scale * len(self.epilog_warp_id), + ), + tx_count=self.num_tma_load_bytes_scale, + cta_layout_vmnk=cluster_layout_vmnk, + tidx=tidx, + mcast_mode_mn=( + 1, + 0, + ), # multicast for scale_a will only happen on the M-mode + defer_sync=True, + ) + # Initialize transform2mma pipeline, which tracks the dependencies between the type conversion + # of A and MMA's consumption on converted A + cta_v_size = cute.size(cluster_layout_vmnk, mode=[0]) + trans2mma_pipeline = pipeline.PipelineAsyncUmma.create( + barrier_storage=storage.a_trans2mma_full_mbar_ptr.data_ptr(), + num_stages=self.num_trans2mma_stage, + producer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, + 32 * len(self.transform_warp_id) * cta_v_size, + ), + consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ) + # Initialize pipeline for tensor B load to MMA + # MMA warp informs TMA warp to proceed to load next tile of B tensor + b_load2mma_pipeline = pipeline.PipelineTmaUmma.create( + barrier_storage=storage.b_load2mma_full_mbar_ptr.data_ptr(), + num_stages=self.num_load2trans_stage, + producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), + consumer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, self.num_mcast_ctas_b + ), + tx_count=self.num_tma_load_bytes_b, + cta_layout_vmnk=cluster_layout_vmnk, + mcast_mode_mn=(0, 1), # multicast for B will only happen on the N-mode + defer_sync=True, + ) + # Initialize accumulator pipeline, which tracks the dependencies between + # MMA's computation of accumulators and epilogue warps' consumption of accumulators + acc_pipeline = pipeline.PipelineUmmaAsync.create( + barrier_storage=storage.acc_full_mbar_ptr.data_ptr(), + num_stages=self.num_acc_stage, + producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), + consumer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, cta_v_size * len(self.epilog_warp_id) + ), + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ) + # Initialize tile info pipeline, which tracks the dependencies between + # tile scheduling warp and other warps + # Skip scheduler warp when computing consumer thread count + num_tile_info_pipeline_consumer_threads = self.threads_per_cta - 32 + tile_info_pipeline = pipeline.PipelineAsync.create( + barrier_storage=storage.tile_info_full_mbar_ptr.data_ptr(), + num_stages=self.num_tile_info_stage, + producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32), + consumer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, + num_tile_info_pipeline_consumer_threads, + ), + defer_sync=True, + ) + + # Tensor memory dealloc barrier init + tmem = utils.TmemAllocator( + storage.tmem_holding_buf, + barrier_for_retrieve=self.tmem_ptr_sync_barrier, + allocator_warp_id=self.epilog_warp_id[0], + is_two_cta=use_2cta_instrs, + two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar_ptr, + ) + + # Cluster arrive after barrier init + pipeline_init_arrive(cluster_shape_mn=self.cluster_shape_mn, is_relaxed=True) + + # Setup smem tensor A/scale/B/C + sC = smem.allocate_tensor( + element_type=self.c_dtype, + layout=c_smem_layout_staged.outer, + byte_alignment=self.smem_buffer_align_bytes, + swizzle=c_smem_layout_staged.inner, + ) + sA_input = smem.allocate_tensor( + element_type=self.a_dtype, + layout=a_smem_layout.outer, + byte_alignment=self.smem_buffer_align_bytes, + swizzle=a_smem_layout.inner, + ) + sS_input = smem.allocate_tensor( + element_type=self.mma_dtype, + layout=scale_smem_layout.outer, + byte_alignment=self.smem_buffer_align_bytes, + swizzle=scale_smem_layout.inner, + ) + sB_input = smem.allocate_tensor( + element_type=self.b_dtype, + layout=b_smem_layout.outer, + byte_alignment=self.smem_buffer_align_bytes, + swizzle=b_smem_layout.inner, + ) + sA_transform = None + # Get smem tensor for transformed A when transform_a_source is SMEM + if cutlass.const_expr(self.transform_a_source == tcgen05.OperandSource.SMEM): + sA_transform = smem.allocate_tensor( + element_type=self.mma_dtype, + layout=a_smem_layout_transform.outer, + byte_alignment=self.smem_buffer_align_bytes, + swizzle=a_smem_layout_transform.inner, + ) + sTile_info = storage.tile_info.get_tensor( + cute.make_layout((4, self.num_tile_info_stage), stride=(1, 4)) + ) + + # Compute multicast mask for A/B buffer full + a_full_mcast_mask = None + b_full_mcast_mask = None + s_full_mcast_mask = None + if cutlass.const_expr(self.is_a_mcast or self.is_b_mcast or use_2cta_instrs): + a_full_mcast_mask = cpasync.create_tma_multicast_mask( + cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=2 + ) + # Scale tensor shares the same multicast mask as the A tensor + s_full_mcast_mask = a_full_mcast_mask + b_full_mcast_mask = cpasync.create_tma_multicast_mask( + cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=1 + ) + + # local_tile partition global tensors + # (bM, bK, loopM, loopK, loopL) + gA_mkl = cute.local_tile( + mA_mkl, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None) + ) + # (bM, bK, loopM, loopK, loopL) + gS_mkl = cute.local_tile( + mS_mkl, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None) + ) + # (bN, bK, loopN, loopK, loopL) + gB_nkl = cute.local_tile( + mB_nkl, cute.slice_(self.mma_tiler, (0, None, None)), (None, None, None) + ) + # (bM, bN, loopM, loopN, loopL) + gC_mnl = cute.local_tile( + mC_mnl, cute.slice_(self.mma_tiler, (None, None, 0)), (None, None, None) + ) + gC_mnl_simt = cute.local_tile( + tensor_c, cute.slice_(self.mma_tiler, (None, None, 0)), (None, None, None) + ) + k_tile_cnt = cute.size(gA_mkl, mode=[3]) + + # Partition global tensor for TiledMMA_A/B/C + thr_mma = tiled_mma.get_slice(mma_tile_coord_v) + # (MMA, MMA_M, MMA_K, loopM, loopK, loopL) + tCgA = thr_mma.partition_A(gA_mkl) + # (MMA, MMA_M, MMA_K, loopM, loopK, loopL) + tCgS = thr_mma.partition_A(gS_mkl) + # (MMA, MMA_N, MMA_K, loopN, loopK, loopL) + tCgB = thr_mma.partition_B(gB_nkl) + # (MMA, MMA_M, MMA_N, loopM, loopN, loopL) + tCgC = thr_mma.partition_C(gC_mnl) + tCgC_simt = thr_mma.partition_C(gC_mnl_simt) + + # Setup copy atom to load A from shared memory for further transformation + copy_atom_a_input = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), self.a_dtype, num_bits_per_copy=32 + ) + a_smem_shape = tiled_mma.partition_shape_A( + cute.dice(self.mma_tiler, (1, None, 1)) + ) + # Setup copy atom to store transformed A into tensor memory or shared memory + copy_atom_a_transform = mixed_input_utils.get_copy_atom_a_transform( + self.mma_dtype, + self.use_2cta_instrs, + self.transform_a_source, + a_smem_shape, + self.a_dtype, + ) + + # Partition global/shared tensor for TMA load A/B + # TMA load A partition_S/D + a_cta_layout = cute.make_layout( + cute.slice_(cluster_layout_vmnk, (0, 0, None, 0)).shape + ) + # ((atom_v, rest_v), STAGE) + # ((atom_v, rest_v), loopM, loopK, loopL) + tAsA, tAgA = cpasync.tma_partition( + tma_atom_a, + block_in_cluster_coord_vmnk[2], + a_cta_layout, + cute.group_modes(sA_input, 0, 3), + cute.group_modes(tCgA, 0, 3), + ) + + thr_mma_leader_cta = tiled_mma.get_slice(0) + # (MMA, MMA_M, MMA_K, STAGE) + tCsS = thr_mma_leader_cta.partition_A(sS_input) + # ((atom_v, rest_v), STAGE) + # ((atom_v, rest_v), loopM, loopK, loopL) + tSsS, tSgS = mixed_input_utils.scale_tma_partition( + tCsS, + tCgS, + tma_atom_s, + block_in_cluster_coord_vmnk, + a_cta_layout, + ) + + # TMA load B partition_S/D + b_cta_layout = cute.make_layout( + cute.slice_(cluster_layout_vmnk, (0, None, 0, 0)).shape + ) + # ((atom_v, rest_v), STAGE) + # ((atom_v, rest_v), loopM, loopK, loopL) + tBsB, tBgB = cpasync.tma_partition( + tma_atom_b, + block_in_cluster_coord_vmnk[1], + b_cta_layout, + cute.group_modes(sB_input, 0, 3), + cute.group_modes(tCgB, 0, 3), + ) + + # (MMA, MMA_N, MMA_K, STAGE) + tCrB = tiled_mma.make_fragment_B(sB_input) + # (MMA, MMA_M, MMA_N) + acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2]) + tCtAcc_fake = tiled_mma.make_fragment_C( + cute.append(acc_shape, self.num_acc_stage) + ) + + # Cluster wait before TMEM alloc and ensure pipelines are ready + pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mn) + + # TMEM allocation + tmem.allocate(self.num_tmem_alloc_cols) + tmem.wait_for_alloc() + # Get the pointer to the TMEM buffer + tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + accumulators = cute.make_tensor(tmem_ptr, tCtAcc_fake.layout) + + tCrA = None + if cutlass.const_expr(self.transform_a_source == tcgen05.OperandSource.TMEM): + tmem_ptr_transform = cute.recast_ptr( + accumulators.iterator + self.num_acc_tmem_cols, dtype=self.mma_dtype + ) + tCrA = cute.make_tensor( + tmem_ptr_transform, + tiled_mma.make_fragment_A(a_smem_layout_transform.outer).layout, + ) + else: + tCrA = tiled_mma.make_fragment_A(sA_transform) + + # Schedule warp + if warp_idx == self.schedule_warp_id: + cute.arch.setmaxregister_decrease(self.num_regs_schedule_warp) + # Persistent tile scheduling loop + tile_sched = utils.StaticPersistentRuntimeTileScheduler.create( + tile_sched_params, + (bidx, bidy, bidz), + cute.arch.grid_dim(), + inner_mode=0, + ) + work_tile = tile_sched.initial_work_tile_info() + tile_info_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.num_tile_info_stage + ) + # Create initial group search state + search_state = ( + mixed_input_utils.create_initial_contiguous_group_search_state() + ) + not_last_tile = cutlass.Boolean(1) + while not_last_tile: + tile_info_pipeline.producer_acquire(tile_info_producer_state) + cluster_tile_coord_mnl = work_tile.tile_idx + cta_tile_coord_m = ( + cluster_tile_coord_mnl[0] * self.cluster_shape_mn[0] + + block_in_cluster_coord_vmnk[1] * cute.size(tiled_mma.thr_id.shape) + + block_in_cluster_coord_vmnk[0] + ) + cta_tile_offset_n = block_in_cluster_coord_vmnk[2] + search_state = mixed_input_utils.contiguous_group_search( + self.cluster_tile_shape_mnk, + group_count, + cluster_tile_coord_mnl[1], + search_state, + cumsum, + 1, # mode index to perform the search. 0 for M and 1 for N + ) + cur_sTile_info = sTile_info[(None, tile_info_producer_state.index)] + not_last_tile = search_state.cur_group_idx <= group_count + # Store tile info into shared memory buffer + with cute.arch.elect_one(): + cur_sTile_info[0] = cta_tile_coord_m + cur_sTile_info[1] = ( + search_state.cur_start + + cta_tile_offset_n * self.cta_tile_shape_mnk[1] + ) + cur_sTile_info[2] = search_state.cur_group_idx - 1 + cur_sTile_info[3] = ( + search_state.cur_boundary + - search_state.cur_start + - (cta_tile_offset_n * self.cta_tile_shape_mnk[1]) + ) + # Fence and barrier to ensure tile info store has finished + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) + self.sched_sync_barrier.arrive_and_wait() + # Commit tile info pipeline + tile_info_pipeline.producer_commit(tile_info_producer_state) + # Advance to next tile + tile_info_producer_state.advance() + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + tile_info_pipeline.producer_tail(tile_info_producer_state) + + # Specialized TMA load warp for A/B tensor + if warp_idx == self.tma_warp_id: + cute.arch.setmaxregister_decrease(self.num_regs_tma_warps) + # Persistent tile scheduling loop + tile_info_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_tile_info_stage + ) + tile_info_pipeline.consumer_wait(tile_info_consumer_state) + work_tile = mixed_input_utils.make_contiguous_group_work_tile_info( + group_count, sTile_info[(None, tile_info_consumer_state.index)] + ) + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) + tile_info_pipeline.consumer_release(tile_info_consumer_state) + tile_info_consumer_state.advance() + a_load2trans_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.num_load2trans_stage + ) + b_load2mma_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.num_load2trans_stage + ) + + while work_tile.is_valid_tile: + tAgA_slice = tAgA[ + ( + None, + work_tile.cta_coord_m // cute.size(tiled_mma.thr_id.shape), + None, + work_tile.group_idx, + ) + ] + # Apply offset to B tensor based on group search result + coord_n_offset = ( + (work_tile.coord_n, 0, 0) + if cutlass.const_expr( + self.b_major_mode == tcgen05.OperandMajorMode.MN + ) + else (0, work_tile.coord_n, 0) + ) + tBgB_slice = cute.make_tensor( + ( + tBgB.iterator[0] + coord_n_offset[0], + coord_n_offset[1] + tBgB.iterator[1], + coord_n_offset[2] + tBgB.iterator[2], + ), + cute.slice_(tBgB.layout, (None, 0, None, 0)), + ) + + a_load2trans_producer_state.reset_count() + peek_load2trans_empty_status = cutlass.Boolean(1) + if a_load2trans_producer_state.count < k_tile_cnt: + peek_load2trans_empty_status = ( + a_load2trans_pipeline.producer_try_acquire( + a_load2trans_producer_state + ) + ) + b_load2mma_producer_state.reset_count() + for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1): + a_load2trans_pipeline.producer_acquire( + a_load2trans_producer_state, peek_load2trans_empty_status + ) + b_load2mma_pipeline.producer_acquire(b_load2mma_producer_state) + # TMA load A/B + cute.copy( + tma_atom_a, + tAgA_slice[(None, a_load2trans_producer_state.count)], + tAsA[(None, a_load2trans_producer_state.index)], + tma_bar_ptr=a_load2trans_pipeline.producer_get_barrier( + a_load2trans_producer_state + ), + mcast_mask=a_full_mcast_mask, + ) + cute.copy( + tma_atom_b, + tBgB_slice[(None, b_load2mma_producer_state.count)], + tBsB[(None, b_load2mma_producer_state.index)], + tma_bar_ptr=b_load2mma_pipeline.producer_get_barrier( + b_load2mma_producer_state + ), + mcast_mask=b_full_mcast_mask, + ) + a_load2trans_pipeline.producer_commit(a_load2trans_producer_state) + b_load2mma_pipeline.producer_commit(b_load2mma_producer_state) + a_load2trans_producer_state.advance() + b_load2mma_producer_state.advance() + if a_load2trans_producer_state.count < k_tile_cnt: + peek_load2trans_empty_status = ( + a_load2trans_pipeline.producer_try_acquire( + a_load2trans_producer_state + ) + ) + # Advance to next tile + tile_info_pipeline.consumer_wait(tile_info_consumer_state) + work_tile = mixed_input_utils.make_contiguous_group_work_tile_info( + group_count, sTile_info[(None, tile_info_consumer_state.index)] + ) + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) + tile_info_pipeline.consumer_release(tile_info_consumer_state) + tile_info_consumer_state.advance() + # Wait A/B buffer empty + a_load2trans_pipeline.producer_tail(a_load2trans_producer_state) + b_load2mma_pipeline.producer_tail(b_load2mma_producer_state) + + # Specialized TMA load for scale tensor + if warp_idx == self.scale_tma_warp_id: + cute.arch.setmaxregister_decrease(self.num_regs_tma_warps) + if cutlass.const_expr(self.scale_mode == TransformMode.ConvertScale): + # Persistent tile scheduling loop + tile_info_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_tile_info_stage + ) + tile_info_pipeline.consumer_wait(tile_info_consumer_state) + work_tile = mixed_input_utils.make_contiguous_group_work_tile_info( + group_count, sTile_info[(None, tile_info_consumer_state.index)] + ) + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) + tile_info_pipeline.consumer_release(tile_info_consumer_state) + tile_info_consumer_state.advance() + scale_load2accu_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.num_scale_load2accu_stage + ) + scale_k_tile_cnt = cute.size(mS_mkl.layout.shape[1][1]) + + while work_tile.is_valid_tile: + # ((atom_v, rest_v), RestK) + tSgS_slice = tSgS[ + ( + None, + work_tile.cta_coord_m // cute.size(tiled_mma.thr_id.shape), + None, + work_tile.group_idx, + ) + ] + # Filter zeros in rest mode + rest_filtered = cute.filter_zeros(tSgS_slice[(0, None)].layout) + tSgS_slice_filtered = cute.make_tensor( + tSgS_slice.iterator, + cute.make_layout( + (tSgS_slice.layout[0].shape, rest_filtered.shape), + stride=(tSgS_slice.layout[0].stride, rest_filtered.stride), + ), + ) + + scale_load2accu_producer_state.reset_count() + peek_scale_load2accu_empty_status = cutlass.Boolean(1) + if scale_load2accu_producer_state.count < scale_k_tile_cnt: + peek_scale_load2accu_empty_status = ( + scale_load2accu_pipeline.producer_try_acquire( + scale_load2accu_producer_state + ) + ) + for k_tile in cutlass.range(0, scale_k_tile_cnt, 1, unroll=1): + scale_load2accu_pipeline.producer_acquire( + scale_load2accu_producer_state, + peek_scale_load2accu_empty_status, + ) + # TMA load scale + cute.copy( + tma_atom_s, + tSgS_slice_filtered[ + (None, scale_load2accu_producer_state.count) + ], + tSsS[(None, scale_load2accu_producer_state.index)], + tma_bar_ptr=scale_load2accu_pipeline.producer_get_barrier( + scale_load2accu_producer_state + ), + mcast_mask=s_full_mcast_mask, + ) + + scale_load2accu_producer_state.advance() + peek_scale_load2accu_empty_status = cutlass.Boolean(1) + if scale_load2accu_producer_state.count < scale_k_tile_cnt: + peek_scale_load2accu_empty_status = ( + scale_load2accu_pipeline.producer_try_acquire( + scale_load2accu_producer_state + ) + ) + # Advance to next tile + tile_info_pipeline.consumer_wait(tile_info_consumer_state) + work_tile = mixed_input_utils.make_contiguous_group_work_tile_info( + group_count, sTile_info[(None, tile_info_consumer_state.index)] + ) + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) + tile_info_pipeline.consumer_release(tile_info_consumer_state) + tile_info_consumer_state.advance() + # Wait scale buffer empty + scale_load2accu_pipeline.producer_tail(scale_load2accu_producer_state) + + # Specialized transform warps + if warp_idx >= self.transform_warp_id[0]: + cute.arch.setmaxregister_increase(self.num_regs_transform_warps) + transform_local_tidx = tidx - 32 * self.transform_warp_id[0] + # Partition tensors for transform input and output and set up the copy atom + # used for loading and storing transformed A tensor + src_copy_a, dst_copy_a, tAsA_input, tAsA_transform = ( + mixed_input_utils.transform_partition( + self.transform_a_source, + self.scale_mode, + copy_atom_a_input, + copy_atom_a_transform, + sA_input, + ( + tCrA + if self.transform_a_source == tcgen05.OperandSource.TMEM + else sA_transform + ), + transform_local_tidx, + ) + ) + # make fragment for input A and transformed A + tArA_load = cute.make_rmem_tensor( + cute.append( + tAsA_input[(None, None, None, None, 0)].shape, + self.transform_k_tile_unroll_factor, + ), + tAsA_input.element_type, + ) + + tArA_transform = cute.make_rmem_tensor( + cute.append( + tAsA_input[(None, None, None, None, 0)].shape, + self.transform_k_tile_unroll_factor, + ), + self.mma_dtype, + ) + # Deduce a sub-tile size and tile tensors + transform_tiler_size = min( + cute.size(cute.coalesce(tAsA_input.layout), mode=[0]), 32 + ) + transform_tiler = cute.make_layout(transform_tiler_size) + + tile_info_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_tile_info_stage + ) + tile_info_pipeline.consumer_wait(tile_info_consumer_state) + work_tile = mixed_input_utils.make_contiguous_group_work_tile_info( + group_count, sTile_info[(None, tile_info_consumer_state.index)] + ) + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) + tile_info_pipeline.consumer_release(tile_info_consumer_state) + tile_info_consumer_state.advance() + a_load2trans_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, + self.num_load2trans_stage, + ) + a_load2trans_consumer_state0 = a_load2trans_consumer_state.clone() + trans2mma_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, + self.num_trans2mma_stage, + ) + trans2mma_producer_state0 = trans2mma_producer_state.clone() + k_tile_cnt_unrolled2 = k_tile_cnt // self.transform_k_tile_unroll_factor + is_tile_cnt_odd = k_tile_cnt % self.transform_k_tile_unroll_factor == 1 + while work_tile.is_valid_tile: + a_load2trans_consumer_state.reset_count() + a_load2trans_consumer_state0, a_load2trans_consumer_state = ( + self.pipeline_state_clone_and_advance(a_load2trans_consumer_state) + ) + trans2mma_producer_state.reset_count() + trans2mma_producer_state0, trans2mma_producer_state = ( + self.pipeline_state_clone_and_advance(trans2mma_producer_state) + ) + for k_tile in cutlass.range(0, k_tile_cnt_unrolled2, 1, unroll=1): + tAsA_input_slice0, tAsA_input_slice1 = ( + self.slice_and_divide_with_index_pair( + tAsA_input, + ( + a_load2trans_consumer_state0.index, + a_load2trans_consumer_state.index, + ), + transform_tiler, + ) + ) + # reg buffer0 and buffer1 for A_load + tArA_load_slice0, tArA_load_slice1 = ( + self.slice_and_divide_with_index_pair( + tArA_load, + (0, 1), + transform_tiler, + ) + ) + # reg buffer0 and buffer1 for A_transform + tArA_transform_buffer0 = tArA_transform[(None, None, None, None, 0)] + tArA_transform_buffer1 = tArA_transform[(None, None, None, None, 1)] + tArA_transform_slice0 = self.divide_tensor_by_tiler( + tArA_transform_buffer0, + transform_tiler, + ) + tArA_transform_slice1 = self.divide_tensor_by_tiler( + tArA_transform_buffer1, + transform_tiler, + ) + # Check if input A data are ready + a_load2trans_pipeline.consumer_wait(a_load2trans_consumer_state0) + a_load2trans_pipeline.consumer_wait(a_load2trans_consumer_state) + trans2mma_pipeline.producer_acquire(trans2mma_producer_state) + # Transformation in buffer0 + for idx in cutlass.range_constexpr( + cute.size(tArA_load_slice0, mode=[1]) + ): + # Load A from shared memory + cute.autovec_copy( + tAsA_input_slice0[(None, idx)], + tArA_load_slice0[(None, idx)], + ) + # Convert it to mma dtype + tensor_transformed = mixed_input_utils.cvt_tensor_a( + tArA_load_slice0[(None, idx)], + self.mma_dtype, + self.shuffle_a, + ) + # Load A from shared memory + cute.autovec_copy( + tAsA_input_slice1[(None, idx)], + tArA_load_slice1[(None, idx)], + ) + tArA_transform_slice0[(None, idx)].store(tensor_transformed) + # Store transformed A to tensor memory or shared memory + mixed_input_utils.store_transformed_a( + tArA_transform_buffer0, + tAsA_transform[ + (None, None, None, None, trans2mma_producer_state0.index) + ], + dst_copy_a, + ) + if cutlass.const_expr( + self.transform_a_source == tcgen05.OperandSource.TMEM + ): + cute.arch.fence_view_async_tmem_store() + else: + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) + # Signal the completion of transformation in buffer0 + trans2mma_pipeline.producer_commit(trans2mma_producer_state0) + a_load2trans_pipeline.consumer_release(a_load2trans_consumer_state0) + for idx in cutlass.range_constexpr( + cute.size(tArA_load_slice1, mode=[1]) + ): + # Convert it to mma dtype + tensor_transformed = mixed_input_utils.cvt_tensor_a( + tArA_load_slice1[(None, idx)], + self.mma_dtype, + self.shuffle_a, + ) + tArA_transform_slice1[(None, idx)].store(tensor_transformed) + # Store transformed A to tensor memory or shared memory + mixed_input_utils.store_transformed_a( + tArA_transform_buffer1, + tAsA_transform[ + (None, None, None, None, trans2mma_producer_state.index) + ], + dst_copy_a, + ) + if cutlass.const_expr( + self.transform_a_source == tcgen05.OperandSource.TMEM + ): + cute.arch.fence_view_async_tmem_store() + else: + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) + # Signal the completion of transformation + trans2mma_pipeline.producer_commit(trans2mma_producer_state) + trans2mma_producer_state.advance() + trans2mma_producer_state0, trans2mma_producer_state = ( + self.pipeline_state_clone_and_advance(trans2mma_producer_state) + ) + + a_load2trans_pipeline.consumer_release(a_load2trans_consumer_state) + a_load2trans_consumer_state.advance() + a_load2trans_consumer_state0, a_load2trans_consumer_state = ( + self.pipeline_state_clone_and_advance( + a_load2trans_consumer_state + ) + ) + # Handle the last tile if needed + if is_tile_cnt_odd: + tAsA_input_slice = tAsA_input[ + (None, None, None, None, a_load2trans_consumer_state0.index) + ] + tAsA_input_slice = self.divide_tensor_by_tiler( + tAsA_input_slice, transform_tiler + ) + tArA_load_slice = tArA_load[(None, None, None, None, 0)] + tArA_load_slice = self.divide_tensor_by_tiler( + tArA_load_slice, transform_tiler + ) + tArA_transform_buffer = tArA_transform[(None, None, None, None, 0)] + tArA_transform_slice = self.divide_tensor_by_tiler( + tArA_transform_buffer, transform_tiler + ) + a_load2trans_pipeline.consumer_wait(a_load2trans_consumer_state0) + trans2mma_pipeline.producer_acquire(trans2mma_producer_state0) + for idx in cutlass.range_constexpr( + cute.size(tArA_load_slice, mode=[1]) + ): + # Load A from shared memory + cute.autovec_copy( + tAsA_input_slice[(None, idx)], + tArA_load_slice[(None, idx)], + ) + # Convert it to mma dtype + tensor_transformed = mixed_input_utils.cvt_tensor_a( + tArA_load_slice[(None, idx)], + self.mma_dtype, + self.shuffle_a, + ) + tArA_transform_slice[(None, idx)].store(tensor_transformed) + a_load2trans_pipeline.consumer_release(a_load2trans_consumer_state0) + # Store transformed A to tensor memory or shared memory + mixed_input_utils.store_transformed_a( + tArA_transform_buffer, + tAsA_transform[ + (None, None, None, None, trans2mma_producer_state0.index) + ], + dst_copy_a, + ) + if cutlass.const_expr( + self.transform_a_source == tcgen05.OperandSource.TMEM + ): + cute.arch.fence_view_async_tmem_store() + else: + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) + # Signal the completion of transformation + trans2mma_pipeline.producer_commit(trans2mma_producer_state0) + trans2mma_producer_state0.advance() + a_load2trans_consumer_state0.advance() + # Keep pipeline state ready for next available buffer + trans2mma_producer_state = trans2mma_producer_state0.clone() + a_load2trans_consumer_state = a_load2trans_consumer_state0.clone() + # Advance to next tile + tile_info_pipeline.consumer_wait(tile_info_consumer_state) + work_tile = mixed_input_utils.make_contiguous_group_work_tile_info( + group_count, sTile_info[(None, tile_info_consumer_state.index)] + ) + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) + tile_info_pipeline.consumer_release(tile_info_consumer_state) + tile_info_consumer_state.advance() + # Wait a_transform buffer empty + trans2mma_pipeline.producer_tail(trans2mma_producer_state0) + + # Specialized MMA warp + if warp_idx == self.mma_warp_id: + cute.arch.setmaxregister_decrease(self.num_regs_mma_warp) + tCtAcc_base = accumulators + # Persistent tile scheduling loop + tile_info_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_tile_info_stage + ) + tile_info_pipeline.consumer_wait(tile_info_consumer_state) + work_tile = mixed_input_utils.make_contiguous_group_work_tile_info( + group_count, sTile_info[(None, tile_info_consumer_state.index)] + ) + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) + tile_info_pipeline.consumer_release(tile_info_consumer_state) + tile_info_consumer_state.advance() + trans2mma_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_trans2mma_stage + ) + b_load2mma_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_load2trans_stage + ) + acc_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.num_acc_stage + ) + while work_tile.is_valid_tile: + b_load2mma_consumer_state.reset_count() + trans2mma_consumer_state.reset_count() + peek_trans2mma_full_status = cutlass.Boolean(1) + if is_leader_cta: + if trans2mma_consumer_state.count < k_tile_cnt: + peek_trans2mma_full_status = ( + trans2mma_pipeline.consumer_try_wait( + trans2mma_consumer_state + ) + ) + k_block_cnt = cute.ceil_div(k_tile_cnt, num_k_tiles_per_scale) + num_k_tiles_executed = 0 + # Loop over K blocks with different scales and commit + # before starting the next tile requiring new scales + for k_block in cutlass.range(0, k_block_cnt, 1, unroll=1): + # (MMA, MMA_M, MMA_N) + tCtAcc = tCtAcc_base[ + (None, None, None, acc_producer_state.index) + ] + acc_pipeline.producer_acquire(acc_producer_state) + cur_num_k_tiles = min( + num_k_tiles_per_scale, k_tile_cnt - num_k_tiles_executed + ) + tiled_mma.set(tcgen05.Field.ACCUMULATE, False) + for k_tile in cutlass.range(0, cur_num_k_tiles, 1, unroll=1): + trans2mma_pipeline.consumer_wait( + trans2mma_consumer_state, peek_trans2mma_full_status + ) + b_load2mma_pipeline.consumer_wait(b_load2mma_consumer_state) + num_kslices = cute.size(tCrA, mode=[2]) + for kslice_idx in cutlass.range( + num_kslices, unroll_full=True + ): + kblock_coord_a = ( + None, + None, + kslice_idx, + trans2mma_consumer_state.index, + ) + kblock_coord_b = ( + None, + None, + kslice_idx, + b_load2mma_consumer_state.index, + ) + + cute.gemm( + tiled_mma, + tCtAcc, + tCrA[kblock_coord_a], + tCrB[kblock_coord_b], + tCtAcc, + ) + # Enable accumulate on tCtAcc after first kblock + tiled_mma.set(tcgen05.Field.ACCUMULATE, True) + trans2mma_pipeline.consumer_release( + trans2mma_consumer_state + ) + b_load2mma_pipeline.consumer_release( + b_load2mma_consumer_state + ) + trans2mma_consumer_state.advance() + b_load2mma_consumer_state.advance() + peek_trans2mma_full_status = cutlass.Boolean(1) + if trans2mma_consumer_state.count < k_tile_cnt: + peek_trans2mma_full_status = ( + trans2mma_pipeline.consumer_try_wait( + trans2mma_consumer_state + ) + ) + num_k_tiles_executed += cur_num_k_tiles + # Async arrive accumulator buffer full + acc_pipeline.producer_commit(acc_producer_state) + acc_producer_state.advance() + + # Advance to next tile + tile_info_pipeline.consumer_wait(tile_info_consumer_state) + work_tile = mixed_input_utils.make_contiguous_group_work_tile_info( + group_count, sTile_info[(None, tile_info_consumer_state.index)] + ) + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) + tile_info_pipeline.consumer_release(tile_info_consumer_state) + tile_info_consumer_state.advance() + # Wait for accumulator buffer empty + acc_pipeline.producer_tail(acc_producer_state) + + # Specialized acc update and epilogue warps + if warp_idx < self.mma_warp_id: + cute.arch.setmaxregister_increase(self.num_regs_epilogue_warps) + epi_tidx = tidx + tCtAcc_base = accumulators + # Construct scale tensor view as C + scale_view_as_C_layout = cute.make_layout( + ( + scale_smem_layout.outer[0].shape, + self.cta_tile_shape_mnk[1], + scale_smem_layout.outer[2].shape, + ), + stride=( + scale_smem_layout.outer[0].stride, + 0, + scale_smem_layout.outer[2].stride, + ), + ) + scale_view_as_C = cute.make_tensor( + sS_input.iterator, + scale_view_as_C_layout, + ) + # Partition for epilogue and accumulator update + tiled_copy_t2r, tTR_tAcc_base, tTR_rAcc, tTR_rAcc_final, tTR_sScale = ( + self.epilog_and_acc_update_tmem_copy_and_partition( + epi_tidx, + tCtAcc_base, + tCgC, + scale_view_as_C, + epi_tile, + use_2cta_instrs, + ) + ) + + tTR_rC = cute.make_rmem_tensor(tTR_rAcc.shape, self.c_dtype) + tiled_copy_r2s, tRS_rC, tRS_sC = ( + mixed_input_utils.epilog_smem_copy_and_partition( + self.c_layout, + self.c_dtype, + self.acc_dtype, + tiled_copy_t2r, + tTR_rC, + epi_tidx, + sC, + ) + ) + (tma_atom_c, bSG_sC, bSG_gC_partitioned, simt_atom, tTR_gC_partitioned) = ( + mixed_input_utils.epilog_gmem_copy_and_partition( + self.c_dtype, + epi_tidx, + tma_atom_c, + tiled_copy_t2r, + tCgC, + tCgC_simt, + epi_tile, + sC, + ) + ) + + # Predicates + thr_mapping = cute.make_identity_tensor( + (self.cta_tile_shape_mnk[0], self.cta_tile_shape_mnk[1]) + ) + thr_mapping_mn = cute.flat_divide(thr_mapping, epi_tile) + thr_copy_t2r = tiled_copy_t2r.get_slice(epi_tidx) + m_thr_offset = thr_copy_t2r.partition_D(thr_mapping_mn) + m_thr_offset = cute.group_modes(m_thr_offset, 3, cute.rank(m_thr_offset)) + + acc_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_acc_stage + ) + scale_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, + self.num_scale_load2accu_stage, + ) + + c_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + 32 * len(self.epilog_warp_id), + ) + c_pipeline = pipeline.PipelineTmaStore.create( + num_stages=self.num_c_stage, + producer_group=c_producer_group, + ) + # Persistent tile scheduling loop + tile_info_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_tile_info_stage + ) + tile_info_pipeline.consumer_wait(tile_info_consumer_state) + work_tile = mixed_input_utils.make_contiguous_group_work_tile_info( + group_count, sTile_info[(None, tile_info_consumer_state.index)] + ) + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) + tile_info_pipeline.consumer_release(tile_info_consumer_state) + tile_info_consumer_state.advance() + num_prev_subtiles = cutlass.Int32(0) + scale_k_tile_cnt = cute.size(mS_mkl.layout.shape[1][1]) + while work_tile.is_valid_tile: + # perform accumulator update with scales + tTR_rAcc_final.fill(0.0) + tTR_rScale = cute.make_rmem_tensor( + cute.slice_(tTR_sScale, (None, None, None, 0, None, 0)).shape, + self.a_scale_dtype, + ) + scale_consumer_state.reset_count() + peek_scale_full_status = cutlass.Boolean(1) + if scale_consumer_state.count < scale_k_tile_cnt: + peek_scale_full_status = scale_load2accu_pipeline.consumer_try_wait( + scale_consumer_state + ) + acc_consumer_state.reset_count() + peek_acc_full_status = cutlass.Boolean(1) + if acc_consumer_state.count < scale_k_tile_cnt: + peek_acc_full_status = acc_pipeline.consumer_try_wait( + acc_consumer_state + ) + for k_tile in cutlass.range(0, scale_k_tile_cnt, 1, unroll=1): + tTR_tAcc = tTR_tAcc_base[ + (None, None, None, None, None, acc_consumer_state.index) + ] + tTR_sScale_slice = cute.slice_( + tTR_sScale, + (None, None, None, 0, None, scale_consumer_state.index), + ) + scale_load2accu_pipeline.consumer_wait( + scale_consumer_state, peek_scale_full_status + ) + cute.autovec_copy(tTR_sScale_slice, tTR_rScale) + acc_pipeline.consumer_wait(acc_consumer_state, peek_acc_full_status) + tTR_tAcc = cute.group_modes(tTR_tAcc, 3, cute.rank(tTR_tAcc)) + subtile_cnt = cute.size(tTR_tAcc.shape, mode=[3]) + for subtile_idx in cutlass.range_constexpr(subtile_cnt): + tTR_tAcc_mn = tTR_tAcc[(None, None, None, subtile_idx)] + cute.copy(tiled_copy_t2r, tTR_tAcc_mn, tTR_rAcc) + tTR_rAcc_subtile = tTR_rAcc_final[ + (None, None, None, subtile_idx) + ] + tTR_rScale_subtile = tTR_rScale[(None, None, None, subtile_idx)] + acc_vec = tTR_rAcc.load() + final_vec = tTR_rAcc_subtile.load() + scale = tTR_rScale_subtile.load().to(self.acc_dtype) + final_vec = acc_vec * scale + final_vec + tTR_rAcc_subtile.store(final_vec) + scale_load2accu_pipeline.consumer_release(scale_consumer_state) + scale_consumer_state.advance() + peek_scale_full_status = cutlass.Boolean(1) + if scale_consumer_state.count < scale_k_tile_cnt: + peek_scale_full_status = ( + scale_load2accu_pipeline.consumer_try_wait( + scale_consumer_state + ) + ) + with cute.arch.elect_one(): + acc_pipeline.consumer_release(acc_consumer_state) + acc_consumer_state.advance() + peek_acc_full_status = cutlass.Boolean(1) + if acc_consumer_state.count < scale_k_tile_cnt: + peek_acc_full_status = acc_pipeline.consumer_try_wait( + acc_consumer_state + ) + # epilogue partition + bSG_gC = bSG_gC_partitioned[ + ( + None, + None, + None, + work_tile.cta_coord_m // cute.size(tiled_mma.thr_id.shape), + 0, + 0, + ) + ] + tma_store_offset_coord = ( + (work_tile.coord_n, 0, 0) + if cutlass.const_expr(self.c_layout.is_n_major_c()) + else (0, work_tile.coord_n, 0) + ) + bSG_gC = cute.make_tensor( + ( + tma_store_offset_coord[0] + bSG_gC.iterator[0], + tma_store_offset_coord[1] + bSG_gC.iterator[1], + tma_store_offset_coord[2] + bSG_gC.iterator[2], + ), + bSG_gC.layout, + ) + tTR_gC = tTR_gC_partitioned[ + ( + None, + None, + None, + None, + None, + work_tile.cta_coord_m // cute.size(tiled_mma.thr_id.shape), + 0, + 0, + ) + ] + tTR_gC = cute.make_tensor( + tTR_gC.iterator + (work_tile.coord_n * tensor_c.layout.stride[1]), + tTR_gC.layout, + ) + bSG_gC = cute.group_modes(bSG_gC, 1, cute.rank(bSG_gC)) + tTR_gC = cute.group_modes(tTR_gC, 3, cute.rank(tTR_gC)) + + # Store accumulator to global memory in subtiles + subtile_cnt = cute.size(tTR_rAcc_final.shape, mode=[3]) + for subtile_idx in cutlass.range(subtile_cnt): + tTR_rAcc_subtile = tTR_rAcc_final[(None, None, None, subtile_idx)] + if work_tile.distance_to_boundary >= self.cta_tile_shape_mnk[1]: + # Convert to C type + acc_vec = tiled_copy_r2s.retile(tTR_rAcc_subtile).load() + acc_vec = acc_vec.to(self.c_dtype) + tRS_rC.store(acc_vec) + num_prev_subtiles += 1 + c_buffer = num_prev_subtiles % self.num_c_stage + # Store C to shared memory + cute.copy( + tiled_copy_r2s, + tRS_rC, + tRS_sC[(None, None, None, c_buffer)], + ) + # Fence and barrier to make sure shared memory store is visible to TMA store + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) + self.epilog_sync_barrier.arrive_and_wait() + # TMA store C to global memory + if warp_idx == self.epilog_warp_id[0]: + cute.copy( + tma_atom_c, + bSG_sC[(None, c_buffer)], + bSG_gC[(None, subtile_idx)], + ) + c_pipeline.producer_commit() + c_pipeline.producer_acquire() + self.epilog_sync_barrier.arrive_and_wait() + else: + # Convert to C type + acc_vec = tTR_rAcc_subtile.load() + acc_vec = acc_vec.to(self.c_dtype) + tTR_rC.store(acc_vec) + # Compute predicate for SIMT store + tCpC = cute.make_rmem_tensor( + cute.make_layout(tTR_rC.shape), + cutlass.Boolean, + ) + m_thr_slice = m_thr_offset[(None, None, None, subtile_idx)] + for i in cutlass.range(cute.size(tCpC), unroll_full=True): + tCpC[i] = ( + m_thr_slice[(i)][0] + + work_tile.cta_coord_m * self.cta_tile_shape_mnk[0] + < tensor_c.shape[0] + ) and (m_thr_slice[(i)][1] < work_tile.distance_to_boundary) + # Store C to global memory + cute.copy( + simt_atom, + cute.flatten(tTR_rC), + cute.flatten(tTR_gC[(None, None, None, subtile_idx)]), + pred=cute.flatten(tCpC), + ) + # Advance to next tile + tile_info_pipeline.consumer_wait(tile_info_consumer_state) + work_tile = mixed_input_utils.make_contiguous_group_work_tile_info( + group_count, sTile_info[(None, tile_info_consumer_state.index)] + ) + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) + tile_info_pipeline.consumer_release(tile_info_consumer_state) + tile_info_consumer_state.advance() + + # Dealloc the tensor memory buffer + tmem.relinquish_alloc_permit() + self.epilog_sync_barrier.arrive_and_wait() + tmem.free(tmem_ptr) + c_pipeline.producer_tail() + + def divide_tensor_by_tiler( + self, tensor: cute.Tensor, transform_tiler: cute.Layout + ) -> cute.Tensor: + """ + Divide the input tensor by given tiler and organize the resulting layout to 2 modes. + The first mode is the tile mode and the second mode is the rest mode. + """ + divided_tensor = cute.flat_divide(tensor, transform_tiler) + divided_tensor = cute.group_modes(divided_tensor, 1, cute.rank(divided_tensor)) + return divided_tensor + + def slice_and_divide_with_index_pair( + self, + tensor: cute.Tensor, + index_pair: tuple[cutlass.Int32, cutlass.Int32], + tiler: cute.Layout, + slice_mode=4, + ) -> tuple[cute.Tensor, cute.Tensor]: + """ + Perform the slice and divide_tensor_by_tiler operation on the sliced tensor with index pair. Coords used for slice are + """ + # pad None before the slice_mode + tensor0_slice, tensor1_slice = ( + tensor[(None,) * slice_mode + (index_pair[0],)], + tensor[(None,) * slice_mode + (index_pair[1],)], + ) + return self.divide_tensor_by_tiler( + tensor0_slice, tiler + ), self.divide_tensor_by_tiler(tensor1_slice, tiler) + + def pipeline_state_clone_and_advance( + self, pipeline_state: pipeline.PipelineState + ) -> tuple[pipeline.PipelineState, pipeline.PipelineState]: + """ + Clones the pipeline state and advances it. + """ + pipeline_state_clone = pipeline_state.clone() + pipeline_state.advance() + return pipeline_state_clone, pipeline_state + + def epilog_and_acc_update_tmem_copy_and_partition( + self, + tidx: cutlass.Int32, + tAcc: cute.Tensor, + gC_mnl: cute.Tensor, + scale_tensor: cute.Tensor, + epi_tile: cute.Tile, + use_2cta_instrs: Union[cutlass.Boolean, bool], + ) -> tuple[cute.TiledCopy, cute.Tensor, cute.Tensor, cute.Tensor, cute.Tensor]: + """ + Partitions source and destination tensors for a tensor memory load together with + the scale and accumulator tensors for the accumulator update. + """ + # Make tiledCopy for tensor memory load + copy_atom_t2r = sm100_utils.get_tmem_load_op( + self.cta_tile_shape_mnk, + self.c_layout, + self.c_dtype, + self.acc_dtype, + epi_tile, + use_2cta_instrs, + ) + # (EPI_TILE_M, EPI_TILE_N, EPI_M, EPI_N, STAGE) + tAcc_epi = cute.flat_divide( + tAcc[((None, None), 0, 0, None)], + epi_tile, + ) + # (EPI_TILE_M, EPI_TILE_N) + tiled_copy_t2r = tcgen05.make_tmem_copy( + copy_atom_t2r, tAcc_epi[(None, None, 0, 0, 0)] + ) + + thr_copy_t2r = tiled_copy_t2r.get_slice(tidx) + # (T2R, T2R_M, T2R_N, EPI_M, EPI_M, STAGE) + tTR_tAcc = thr_copy_t2r.partition_S(tAcc_epi) + + # (EPI_TILE_M, EPI_TILE_N, EPI_M, EPI_N, loopM, loopN, loopL) + gC_mnl_epi = cute.flat_divide( + gC_mnl[((None, None), 0, 0, None, None, None)], epi_tile + ) + sScale_epi = cute.flat_divide(scale_tensor, epi_tile) + # (T2R, T2R_M, T2R_N, EPI_M, EPI_N, loopM, loopN, loopL) + tTR_gC = thr_copy_t2r.partition_D(gC_mnl_epi) + tTR_sScale = thr_copy_t2r.partition_D(sScale_epi) + # (T2R, T2R_M, T2R_N) + tTR_rAcc = cute.make_rmem_tensor( + tTR_gC[(None, None, None, 0, 0, 0, 0, 0)].shape, self.acc_dtype + ) + # (T2R, T2R_M, T2R_N, EPI_M, EPI_N) + tTR_rAcc_final_ = cute.make_rmem_tensor( + tTR_gC[(None, None, None, None, None, 0, 0, 0)].shape, self.acc_dtype + ) + tTR_rAcc_final = cute.group_modes( + tTR_rAcc_final_, 3, cute.rank(tTR_rAcc_final_) + ) + return tiled_copy_t2r, tTR_tAcc, tTR_rAcc, tTR_rAcc_final, tTR_sScale + + @staticmethod + def _compute_stages_and_tmem_cols( + tiled_mma: cute.TiledMma, + mma_tiler_mnk: tuple[int, int, int], + cta_tile_shape_mnk: tuple[int, int, int], + epi_tile: cute.Tile, + a_dtype: type[cutlass.Numeric], + b_dtype: type[cutlass.Numeric], + c_dtype: type[cutlass.Numeric], + c_layout: utils.LayoutEnum, + transform_a_source: tcgen05.OperandSource, + scale_granularity_m: int, + scale_granularity_k: int, + smem_buffer_align_bytes: int, + ) -> tuple[int, int, int, int, int, int, int, int]: + """ + Compute pipeline stages and TMEM column allocation configurations. + """ + # Compute tmem columns required for accumulator + acc_shape = tiled_mma.partition_shape_C(mma_tiler_mnk[:2]) + tCtAcc_stage1 = tiled_mma.make_fragment_C(cute.append(acc_shape, 1)) + num_tmem_acc_col_per_stage = cute.round_up( + tcgen05.find_tmem_tensor_col_offset(tCtAcc_stage1), 2 + ) + # Heuristic to decide the number of stages for accumulator + sm100_tmem_columns = cute.arch.get_max_tmem_alloc_cols("sm_100") + accumulator_stage_count = sm100_tmem_columns // num_tmem_acc_col_per_stage + if transform_a_source == tcgen05.OperandSource.TMEM: + if num_tmem_acc_col_per_stage < 128: + accumulator_stage_count = 3 + elif num_tmem_acc_col_per_stage < 256: + accumulator_stage_count = 2 + else: + accumulator_stage_count = 1 + # transformed A in 16bit, thus 1 tmem column could hold 2 elements + num_elts_per_tmem_col = 32 // tiled_mma.op.a_dtype.width + num_tmem_cols_a_per_stage = cute.round_up( + ( + cta_tile_shape_mnk[2] // num_elts_per_tmem_col + if transform_a_source == tcgen05.OperandSource.TMEM + else 0 + ), + 4, + ) + + bytes_per_pipeline_stage = 16 + # By default, we use 2 stages for tile info + num_tile_info_stage = 2 + tile_info_bytes = ( + cute.size_in_bytes(cute.Int32, cute.make_layout((4, num_tile_info_stage))) + + bytes_per_pipeline_stage * num_tile_info_stage + ) + # Just keep 1 stage as tileN is small for decoding cases + c_stage_count = 1 + c_smem_layout_staged_one = sm100_utils.make_smem_layout_epi( + c_dtype, + c_layout, + epi_tile, + 1, + ) + c_bytes_per_stage = cute.size_in_bytes(c_dtype, c_smem_layout_staged_one) + c_bytes = c_bytes_per_stage * c_stage_count + + smem_capacity = utils.get_smem_capacity_in_bytes("sm_100") + # Ensure we have 4 buffers for scale tiles needed for 1 CTA tile + a_scale_k_mode = max(cta_tile_shape_mnk[2] // scale_granularity_k, 1) + a_scale_m_mode = max(cta_tile_shape_mnk[0] // scale_granularity_m, 1) + scale_load2accu_stage_count = 4 + a_scale_bytes_per_stage = cute.round_up( + cute.size_in_bytes( + tiled_mma.op.a_dtype, + cute.make_layout((a_scale_m_mode, a_scale_k_mode)), + ), + smem_buffer_align_bytes, + ) + a_scale_bytes = ( + a_scale_bytes_per_stage + bytes_per_pipeline_stage + ) * scale_load2accu_stage_count + carveout_smem_bytes = ( + bytes_per_pipeline_stage * accumulator_stage_count + + a_scale_bytes + + c_bytes + + tile_info_bytes + ) + + # Compute transform stages if A is in TMEM + num_tmem_acc_cols = cute.round_up( + accumulator_stage_count * num_tmem_acc_col_per_stage, 4 + ) + + transform2mma_stage_count_a_source_tmem_potential = ( + (sm100_tmem_columns - num_tmem_acc_cols) // num_tmem_cols_a_per_stage + if transform_a_source == tcgen05.OperandSource.TMEM + else -1 + ) + if ( + transform_a_source == tcgen05.OperandSource.TMEM + and transform2mma_stage_count_a_source_tmem_potential <= 0 + ): + raise ValueError("Not enough TMEM capacity for selected tile size") + a_load_bytes_per_stage = cute.round_up( + cute.size_in_bytes( + a_dtype, + cute.make_layout((cta_tile_shape_mnk[0], cta_tile_shape_mnk[2])), + ), + smem_buffer_align_bytes, + ) + b_load_bytes_per_stage = cute.round_up( + cute.size_in_bytes( + b_dtype, + cute.make_layout( + ( + cta_tile_shape_mnk[1] // cute.size(tiled_mma.thr_id), + cta_tile_shape_mnk[2], + ) + ), + ), + smem_buffer_align_bytes, + ) + ab_load_bytes_per_stage = ( + a_load_bytes_per_stage + + b_load_bytes_per_stage + + 2 * bytes_per_pipeline_stage + ) + a_transform_bytes_per_stage = ( + cute.round_up( + cute.size_in_bytes( + tiled_mma.op.a_dtype, + cute.make_layout((cta_tile_shape_mnk[0], cta_tile_shape_mnk[2])), + ), + smem_buffer_align_bytes, + ) + if transform_a_source == tcgen05.OperandSource.SMEM + else 0 + ) + + a_transform_bytes_per_stage = ( + a_transform_bytes_per_stage + bytes_per_pipeline_stage + ) + transform2mma_stage_count_a_source_smem_potential = ( + smem_capacity - carveout_smem_bytes + ) // (ab_load_bytes_per_stage + a_transform_bytes_per_stage) + transform2mma_stage_count = ( + min( + transform2mma_stage_count_a_source_tmem_potential, + transform2mma_stage_count_a_source_smem_potential, + ) + if transform_a_source == tcgen05.OperandSource.TMEM + else transform2mma_stage_count_a_source_smem_potential + ) + load2transform_stage_count = ( + smem_capacity + - carveout_smem_bytes + - (transform2mma_stage_count * a_transform_bytes_per_stage) + ) // ab_load_bytes_per_stage + if ( + load2transform_stage_count < 2 + or transform2mma_stage_count < 2 + or accumulator_stage_count < 1 + ): + raise ValueError("Not enough SMEM or TMEM capacity for selected tile size") + num_tmem_a_cols = transform2mma_stage_count * num_tmem_cols_a_per_stage + # Check if we can increase c_stage_count with leftover smem + c_stage_count += ( + smem_capacity + - load2transform_stage_count * ab_load_bytes_per_stage + - transform2mma_stage_count * a_transform_bytes_per_stage + - scale_load2accu_stage_count * a_scale_bytes_per_stage + - c_bytes + ) // c_bytes_per_stage + + return ( + load2transform_stage_count, + scale_load2accu_stage_count, + transform2mma_stage_count, + accumulator_stage_count, + c_stage_count, + num_tile_info_stage, + num_tmem_acc_cols, + num_tmem_a_cols, + ) + + @staticmethod + def _compute_grid( + c: cute.Tensor, + cta_tile_shape_mnk: tuple[int, int, int], + cluster_shape_mn: tuple[int, int], + max_active_clusters: cutlass.Constexpr, + ) -> tuple[utils.PersistentTileSchedulerParams, tuple[int, int, int]]: + """ + Use persistent tile scheduler to compute the grid size for the output tensor C. + """ + c_shape = cute.slice_(cta_tile_shape_mnk, (None, None, 0)) + gc = cute.zipped_divide(c, tiler=c_shape) + num_ctas_mnl = gc[(0, (None, None, None))].shape + cluster_shape_mnl = (*cluster_shape_mn, 1) + + tile_sched_params = utils.PersistentTileSchedulerParams( + num_ctas_mnl, cluster_shape_mnl + ) + grid = (cluster_shape_mn[0], cluster_shape_mn[1], max_active_clusters) + + return tile_sched_params, grid + + def can_implement( + mnkl: tuple[int, int, int, int], + a_dtype: type[cutlass.Numeric], + b_dtype: type[cutlass.Numeric], + c_dtype: type[cutlass.Numeric], + a_major: str, + b_major: str, + c_major: str, + scale_granularity_m: int, + scale_granularity_k: int, + mma_tiler: tuple[int, int, int], + cluster_shape_mn: tuple[int, int], + use_2cta_instrs: bool, + ) -> bool: + """ + Check if the kernel can be implemented for the given tensor shapes and data types. + """ + m, n, k, l = mnkl + + if not mixed_input_utils.is_valid_mma_tiler_and_cluster_shape( + mma_tiler, cluster_shape_mn, use_2cta_instrs + ): + return False + if not mixed_input_utils.is_valid_scale_granularity( + scale_granularity_m, scale_granularity_k, a_dtype, k, mma_tiler[2] + ): + return False + + # Check tensor alignment + def check_contiguous_NB_alignment( + dtype, contiguous_dim_size, expected_align_bytes + ): + expected_alignment = expected_align_bytes * 8 // dtype.width + return contiguous_dim_size % expected_alignment == 0 + + if not ( + check_contiguous_NB_alignment(a_dtype, m if a_major == "m" else k, 16) + and check_contiguous_NB_alignment(b_dtype, n if b_major == "n" else k, 16) + ): + return False + return True + + +def get_advanced_compiler_control_path(): + """ + Return the path to the advanced compiler control file of this example. If not found, return None. + """ + import os + + need_advanced_compiler_control = False + try: + from cutlass import CUDA_VERSION + + if CUDA_VERSION.major == 13 and CUDA_VERSION.minor == 1: + need_advanced_compiler_control = True + except ImportError: + pass + + if not need_advanced_compiler_control: + return None + # Get the path to the advanced compiler control file + current_dir = os.path.dirname(os.path.abspath(__file__)) + target_path = os.path.join(current_dir, "../../advanced_compiler_control/gemm0.bin") + if os.path.exists(target_path): + print(f"Found advanced compiler control file at {target_path}") + return target_path + else: + return None + + +def run( + mnkl: tuple[int, int, int, int], + scale_granularity_m: int, + scale_granularity_k: int, + a_dtype: type[cutlass.Numeric], + b_dtype: type[cutlass.Numeric], + c_dtype: type[cutlass.Numeric], + acc_dtype: type[cutlass.Numeric], + a_major: str, + b_major: str, + c_major: str, + mma_tiler_mnk: tuple[int, int, int], + cluster_shape_mn: tuple[int, int], + use_2cta_instrs: bool, + tolerance: float, + warmup_iterations: int = 0, + iterations: int = 1, + skip_ref_check: bool = False, + uniform_group_sizes: bool = False, + use_cold_l2: bool = False, + **kwargs, +) -> None: + """ + Run the mixed-input GEMM kernel with specified parameters. + + This function creates tensors, validates parameters, executes the kernel, + optionally compares results with a reference implementation and reports + kernel execution time. + """ + m, n, k, l = mnkl + + if not torch.cuda.is_available(): + raise ValueError("CUDA is not available") + + # Check if given configuration is supported + if not GroupedMixedInputGemmAccScaleKernel.can_implement( + mnkl, + a_dtype, + b_dtype, + c_dtype, + a_major, + b_major, + c_major, + scale_granularity_m, + scale_granularity_k, + mma_tiler_mnk, + cluster_shape_mn, + use_2cta_instrs, + ): + raise ValueError("GEMM configuration not supported") + + # Get current CUDA stream from PyTorch + torch_stream = torch.cuda.current_stream() + # Get the raw stream pointer as a CUstream + current_stream = cuda.CUstream(torch_stream.cuda_stream) + + group_count = l + shuffle_a = mixed_input_utils.is_shuffle_a( + a_major, k, a_dtype, b_dtype, scale_granularity_k + ) + # shuffle is supported since CUDA 13.1 + shuffle_supported = False + try: + from cutlass import CUDA_VERSION + + if CUDA_VERSION.major > 13 or ( + CUDA_VERSION.major == 13 and CUDA_VERSION.minor >= 1 + ): + shuffle_supported = True + except ImportError: + pass + + shuffle_a = shuffle_a and shuffle_supported + mixed_input_gemm = GroupedMixedInputGemmAccScaleKernel( + scale_granularity_m, + scale_granularity_k, + acc_dtype, + use_2cta_instrs, + mma_tiler_mnk, + cluster_shape_mn, + group_count, + shuffle_a, + ) + torch.manual_seed(2025) + ( + a_tensor, + a_scale_tensor, + b_tensor, + cumsum_tensor, + c_tensor, + a_torch_cpu, + a_scale_torch_cpu, + b_torch_cpu, + cumsum_torch_cpu, + c_torch_gpu, + ) = create_tensors( + l, + m, + n, + k, + a_major, + b_major, + c_major, + a_dtype, + b_dtype, + c_dtype, + shuffle_a, + scale_granularity_m, + scale_granularity_k, + uniform_group_sizes, + ) + + max_active_clusters = utils.HardwareInfo().get_max_active_clusters( + cluster_shape_mn[0] * cluster_shape_mn[1], + ) + advanced_compiler_options = None + advanced_compiler_control_path = get_advanced_compiler_control_path() + if advanced_compiler_control_path: + advanced_compiler_options = ( + f"--ptxas-options '--apply-controls={advanced_compiler_control_path}'" + ) + + compiled_kernel = cute.compile( + mixed_input_gemm, + a_tensor, + a_scale_tensor, + b_tensor, + cumsum_tensor, + c_tensor, + max_active_clusters, + current_stream, + options=advanced_compiler_options, + ) + + if not skip_ref_check: + compiled_kernel( + a_tensor, + a_scale_tensor, + b_tensor, + cumsum_tensor, + c_tensor, + current_stream, + ) + run_ref_and_compare( + a_torch_cpu, + b_torch_cpu, + a_scale_torch_cpu, + cumsum_torch_cpu, + c_torch_gpu, + c_dtype, + tolerance, + ) + + # Early return if no performance measurement is needed + if iterations <= 0: + return + + def generate_tensors(): + ( + a_tensor, + a_scale_tensor, + b_tensor, + cumsum_tensor, + c_tensor, + a_torch_cpu, + a_scale_torch_cpu, + b_torch_cpu, + cumsum_torch_cpu, + c_torch_gpu, + ) = create_tensors( + l, + m, + n, + k, + a_major, + b_major, + c_major, + a_dtype, + b_dtype, + c_dtype, + shuffle_a, + scale_granularity_m, + scale_granularity_k, + uniform_group_sizes, + ) + return testing.JitArguments( + a_tensor, a_scale_tensor, b_tensor, cumsum_tensor, c_tensor, current_stream + ) + + workspace_count = 1 + if use_cold_l2: + one_workspace_bytes = ( + a_torch_cpu.numel() * a_torch_cpu.element_size() + + b_torch_cpu.numel() * b_torch_cpu.element_size() + + c_torch_gpu.numel() * c_torch_gpu.element_size() + + a_scale_torch_cpu.numel() * a_scale_torch_cpu.element_size() + if a_scale_torch_cpu is not None + else 0 + ) + workspace_count = testing.get_workspace_count( + one_workspace_bytes, warmup_iterations, iterations + ) + + exec_time = testing.benchmark( + compiled_kernel, + workspace_generator=generate_tensors, + workspace_count=workspace_count, + stream=current_stream, + warmup_iterations=warmup_iterations, + iterations=iterations, + ) + + return exec_time # Return execution time in microseconds + + +if __name__ == "__main__": + + def parse_comma_separated_ints(s: str) -> tuple[int, ...]: + try: + return tuple(int(x.strip()) for x in s.split(",")) + except ValueError: + raise argparse.ArgumentTypeError( + "Invalid format. Expected comma-separated integers." + ) + + parser = argparse.ArgumentParser() + parser.add_argument( + "--mnkl", type=parse_comma_separated_ints, default=(128, 128, 128, 1) + ) + parser.add_argument( + "--mma_tiler_mnk", type=parse_comma_separated_ints, default=(128, 128, 128) + ) + parser.add_argument( + "--cluster_shape_mn", type=parse_comma_separated_ints, default=(1, 1) + ) + parser.add_argument( + "--use_2cta_instrs", + action="store_true", + help="Enable 2CTA MMA instructions feature", + ) + parser.add_argument( + "--a_dtype", + type=cutlass.dtype, + default=cutlass.Int4, + choices=[cutlass.Int8, cutlass.Uint8, cutlass.Int4], + ) + parser.add_argument( + "--b_dtype", + type=cutlass.dtype, + default=cutlass.BFloat16, + choices=[cutlass.BFloat16, cutlass.Float16], + ) + parser.add_argument("--c_dtype", type=cutlass.dtype, default=cutlass.BFloat16) + parser.add_argument("--acc_dtype", type=cutlass.dtype, default=cutlass.Float32) + parser.add_argument("--a_major", choices=["k", "m"], type=str, default="m") + parser.add_argument("--b_major", choices=["k", "n"], type=str, default="k") + parser.add_argument("--c_major", choices=["n", "m"], type=str, default="n") + parser.add_argument( + "--scale_granularity_m", + type=int, + default=1, + help="Scale granularity along M dimension.", + ) + parser.add_argument( + "--scale_granularity_k", + type=int, + default=128, + help="Scale granularity along K dimension.", + ) + parser.add_argument( + "--tolerance", type=float, default=1e-01, help="Tolerance for validation" + ) + parser.add_argument( + "--warmup_iterations", type=int, default=0, help="Warmup iterations" + ) + parser.add_argument( + "--iterations", + type=int, + default=1, + help="Number of iterations to run the kernel", + ) + parser.add_argument( + "--skip_ref_check", action="store_true", help="Skip reference checking" + ) + parser.add_argument( + "--uniform_group_sizes", action="store_true", help="Use uniform group sizes" + ) + args = parser.parse_args() + print(f"skip_ref_check={args.skip_ref_check}") + run( + args.mnkl, + args.scale_granularity_m, + args.scale_granularity_k, + args.a_dtype, + args.b_dtype, + args.c_dtype, + args.acc_dtype, + args.a_major, + args.b_major, + args.c_major, + args.mma_tiler_mnk, + args.cluster_shape_mn, + args.use_2cta_instrs, + args.tolerance, + args.warmup_iterations, + args.iterations, + args.skip_ref_check, + args.uniform_group_sizes, + ) + print("PASS") diff --git a/examples/python/CuTeDSL/blackwell/mixed_input_gemm/mixed_input_gemm.py b/examples/python/CuTeDSL/blackwell/mixed_input_gemm/mixed_input_gemm.py new file mode 100644 index 00000000..2517e69a --- /dev/null +++ b/examples/python/CuTeDSL/blackwell/mixed_input_gemm/mixed_input_gemm.py @@ -0,0 +1,2321 @@ +# 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. + +import argparse +from math import log2, ceil +from typing import Optional, Union +import os +import sys + +import cuda.bindings.driver as cuda + +import cutlass +import cutlass.cute as cute +import cutlass.pipeline as pipeline +from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait +import cutlass.utils as utils +import cutlass.utils.blackwell_helpers as sm100_utils +import cutlass.utils.mixed_input_helpers as mixed_input_utils +from cutlass.utils.mixed_input_helpers import TransformMode +import cutlass.cute.testing as testing +from cutlass.cute.nvgpu import cpasync, tcgen05 + +if __name__ == "__main__": + current_dir = os.path.dirname(os.path.abspath(__file__)) + sys.path.insert(0, os.path.join(current_dir, "../..")) + +from blackwell.mixed_input_gemm.mixed_input_host_utils import ( + create_tensors_for_batched_mixed_input_gemm as create_tensors, + run_batched_mixed_input_ref_and_compare as run_ref_and_compare, +) + +""" +A mixed-input GEMM example for the NVIDIA Blackwell SM100 architecture using CUTE DSL. + +This example demonstrates an implementation of mixed-input GEMM using a TMA plus Blackwell SM100 TensorCore +warp-specialized persistent kernel. + +The inputs A and B have different data types. In this example, it's assumed that A is the narrow-precision tensor +and B holds data with a wider precision. +MMA will work in the wide precision of tensor B and tensor A will be transformed to the wide precision of tensor B +following 1 of the 2 possible modes as follows: + +1. convert-only mode: + C = type_convert(A) x B + +In convert-only mode, tensor A is directly converted to the wide precision of tensor B. + +2. convert-scale mode: + C = (type_convert(A) * scale) x B + +In convert-scale mode, tensor A is first converted to the wide precision of tensor B and then scaled by the scale tensor. +The scale tensor is in the same precision as tensor B. +The mode is determined by tensor A's data type as follows: +- if tensor A is in int8 or uint8, convert-only mode is used. +- if tensor A is in int4, convert-scale mode is used. + +The output tensor C could have the same precision as tensor B or fp32. + +To run this example: + +.. code-block:: bash + + python examples/blackwell/mixed_input_gemm/mixed_input_gemm.py \ + --a_dtype Int8 --b_dtype BFloat16 \ + --scale_granularity_m 0 --scale_granularity_k 0 \ + --c_dtype BFloat16 --acc_dtype Float32 \ + --mma_tiler_mnk 128,128,64 --cluster_shape_mn 1,1 \ + --mnkl 256,512,8192,1 + +Input A and B have int8 and bf16 data types, respectively. The Blackwell tcgen05 MMA tile shape +is specified as (128,128,64) and the cluster shape is (1,1). The MMA accumulator and output data type +are set as fp32 and bf16, respectively. As tensor A is int8, convert-only mode is used. +scale_granularity_m and scale_granularity_k are set as 0 for convert-only mode. + +Here is an example of running convert-scale mode: + +.. code-block:: bash + + python examples/blackwell/mixed_input_gemm/mixed_input_gemm.py \ + --a_dtype Int4 --b_dtype BFloat16 \ + --scale_granularity_m 1 --scale_granularity_k 256 \ + --c_dtype BFloat16 --acc_dtype Float32 \ + --mma_tiler_mnk 256,128,128 --cluster_shape_mn 2,1 \ + --use_2cta_instrs --use_tma_store \ + --mnkl 1024,8192,6144,16 + +Input A and B have int4 and bf16 data types, respectively. The scale granularity is set as (1,256), +which means each element along the m mode of tensor A has its own scale element and 256 contiguous elements +along the k mode share the same scale element. There is no scale reuse along the L mode. If the GEMM shape is +(M, N, K, L), then the scale tensor shape is (M // scale_granularity_m, K // scale_granularity_k, L), +which is (1024, 6144/256, 16) in this example. +The Blackwell tcgen05 MMA tile shape is specified as (256,128,128) and tcgen05 2CTA feature is enabled. +The cluster shape is (2,1). The MMA accumulator and output data type are set as fp32 and bf16, respectively. +As tensor A is int4, the convert-scale mode is used. + +To collect performance with NCU profiler: + +.. code-block:: bash + + ncu python examples/blackwell/mixed_input_gemm/mixed_input_gemm.py \ + --a_dtype Int8 --b_dtype BFloat16 \ + --scale_granularity_m 0 --scale_granularity_k 0 \ + --c_dtype BFloat16 --acc_dtype Float32 \ + --mma_tiler_mnk 128,128,64 --cluster_shape_mn 1,1 \ + --mnkl 256,512,8192,1 \ + --warmup_iterations 1 --iterations 10 --skip_ref_check + +Besides the requirements from the Blackwell dense GEMM example, there are some constraints for this example: +* The narrow-precision is constrained to be int8, uint8, or int4 and the other data type is bf16 or f16. +* Output data types could only be fp16, bf16, or fp32. +* The scale_granularity_m must be 1 currently. +* The scale_granularity_k must be a multiple of mma_tiler_k and also be divisible by gemm_k. +* The scale tensor must be in M-major mode. +* OOB tiles are not allowed when TMA store is disabled +""" + + +class MixedInputGemmKernel: + """ + Mixed-input GEMM kernel for NVIDIA Blackwell SM100 architecture. + + This kernel supports GEMM operations where input tensors A and B have different + data types, with tensor A being transformed to the precision of tensor B before + matrix multiplication. + + :param scale_granularity_m: Number of elements sharing the same scale factor along the M mode + :type scale_granularity_m: int + :param scale_granularity_k: Number of elements sharing the same scale factor along the K mode + :type scale_granularity_k: int + :param acc_dtype: Data type for accumulation during computation + :type acc_dtype: type[cutlass.Numeric] + :param use_2cta_instrs: Whether to use CTA group 2 for advanced thread cooperation + :type use_2cta_instrs: bool + :param mma_tiler_mnk: Shape of the Matrix Multiply-Accumulate (MMA) tile (M, N, K) + :type mma_tiler_mnk: tuple[int, int, int] + :param cluster_shape_mn: Cluster dimensions (M,N) for parallel processing + :type cluster_shape_mn: tuple[int, int] + :param use_tma_store: Whether to use Tensor Memory Access (TMA) for storing results + :type use_tma_store: bool + :param shuffle_a: Whether to use shuffle intrinsic for int4-to-bf16 conversion + :type shuffle_a: bool + """ + + def __init__( + self, + scale_granularity_m: int, + scale_granularity_k: int, + acc_dtype: type[cutlass.Numeric], + use_2cta_instrs: bool, + mma_tiler_mnk: tuple[int, int, int], + cluster_shape_mn: tuple[int, int], + use_tma_store: bool, + shuffle_a: bool, + ): + """ + Initializes the mixed-input GEMM kernel with a specified configuration. + """ + # Scale granularity defines how many elements share the same scale factor + # along the M and K modes. + self.scale_granularity_m = scale_granularity_m + self.scale_granularity_k = scale_granularity_k + # Set transform mode + if cutlass.const_expr( + self.scale_granularity_m == 0 and self.scale_granularity_k == 0 + ): + self.scale_mode = TransformMode.ConvertOnly + else: + self.scale_mode = TransformMode.ConvertScale + self.acc_dtype = acc_dtype + self.use_2cta_instrs = use_2cta_instrs + self.cluster_shape_mn = cluster_shape_mn + self.mma_tiler = mma_tiler_mnk + self.use_tma_store = use_tma_store + self.shuffle_a = shuffle_a + self.cta_group = ( + tcgen05.CtaGroup.TWO if self.use_2cta_instrs else tcgen05.CtaGroup.ONE + ) + # Set specialized warp ids + self.epilog_warp_id = ( + 0, + 1, + 2, + 3, + ) + self.mma_warp_id = 4 + self.tma_warp_id = 5 + self.scale_tma_warp_id = 6 + self.idle_warp_id = 7 + # 4 warps to do the transformation + self.transform_warp_id = ( + 8, + 9, + 10, + 11, + ) + self.num_regs_epilogue_warps = 192 + self.num_regs_mma_warp = 96 + self.num_regs_tma_warps = 96 + self.num_regs_transform_warps = 208 + self.num_regs_idle_warp = 24 + self.threads_per_cta = 32 * ( + max( + ( + self.mma_warp_id, + self.tma_warp_id, + self.scale_tma_warp_id, + *self.epilog_warp_id, + *self.transform_warp_id, + ) + ) + + 1 + ) + + # Set barrier id for epilogue sync, tmem ptr sync, and transform sync + self.epilog_sync_barrier = pipeline.NamedBarrier( + 1, 32 * len(self.epilog_warp_id) + ) + self.tmem_ptr_sync_barrier = pipeline.NamedBarrier(2, self.threads_per_cta) + self.transform_sync_barrier = pipeline.NamedBarrier( + 3, 32 * len(self.transform_warp_id) + ) + + self.smem_buffer_align_bytes = 1024 + + def _setup_attributes(self): + """Set up configurations that are dependent on GEMM inputs + + This method configures various attributes based on the input tensor properties + (data types, leading dimensions) and kernel settings: + - Deduce where the transformed A tensor is stored + - Configuring tiled MMA + - Computing MMA/cluster/tile shapes + - Computing cluster layout + - Computing multicast CTAs for A/B + - Computing epilogue sub-tile + - Setting up A/scale/B/C stage counts in shared memory + - Setting up transformed A stage count in shared memory or tensor memory + - Computing A/transformed A/scale/B/C memory layout + - Computing tensor memory allocation columns + """ + # Deduce where the transformed A tensor is stored, shared memory(SMEM) or tensor memory(TMEM) + self.transform_a_source = mixed_input_utils.get_transform_a_source( + self.a_major_mode + ) + tiled_mma = sm100_utils.make_trivial_tiled_mma( + self.mma_dtype, + self.a_major_mode, + self.b_major_mode, + self.acc_dtype, + self.cta_group, + self.mma_tiler[:2], + self.transform_a_source, + ) + self.cta_tile_shape_mnk = ( + self.mma_tiler[0] // cute.size(tiled_mma.thr_id.shape), + self.mma_tiler[1], + self.mma_tiler[2], + ) + self.cluster_layout_vmnk = cute.tiled_divide( + cute.make_layout((*self.cluster_shape_mn, 1)), + (tiled_mma.thr_id.shape,), + ) + self.num_mcast_ctas_a = cute.size(self.cluster_layout_vmnk.shape[2]) + self.num_mcast_ctas_b = cute.size(self.cluster_layout_vmnk.shape[1]) + self.is_a_mcast = self.num_mcast_ctas_a > 1 + self.is_b_mcast = self.num_mcast_ctas_b > 1 + + if cutlass.const_expr(self.use_tma_store): + self.epi_tile = sm100_utils.compute_epilogue_tile_shape( + self.cta_tile_shape_mnk, + self.use_2cta_instrs, + self.c_layout, + self.c_dtype, + ) + else: + self.epi_tile = self.cta_tile_shape_mnk[:2] + + # Compute tensor memory(TMEM) columns and stages for each pipeline + ( + self.num_load2trans_stage, + self.num_scale_load2trans_stage, + self.num_trans2mma_stage, + self.num_acc_stage, + self.num_c_stage, + self.num_acc_tmem_cols, + self.num_a_tmem_cols, + ) = self._compute_stages_and_tmem_cols( + tiled_mma, + self.mma_tiler, + self.cta_tile_shape_mnk, + self.epi_tile, + self.a_dtype, + self.b_dtype, + self.c_dtype, + self.c_layout, + self.transform_a_source, + self.scale_granularity_m, + self.scale_granularity_k, + self.smem_buffer_align_bytes, + self.use_tma_store, + self.scale_mode, + ) + + # Align TMEM columns for allocation + # TMEM allocation requires power-of-2 column alignment + # and must meet minimum allocation requirements + self.num_tmem_alloc_cols = cute.round_up( + self.num_acc_tmem_cols + self.num_a_tmem_cols, + cute.arch.get_min_tmem_alloc_cols("sm_100"), + ) + self.num_tmem_alloc_cols = 2 ** (ceil(log2(self.num_tmem_alloc_cols))) + # Get smem layout for C tensor when TMA store is enabled + self.c_smem_layout_staged = ( + sm100_utils.make_smem_layout_epi( + self.c_dtype, + self.c_layout, + self.epi_tile, + self.num_c_stage, + ) + if self.use_tma_store + else None + ) + # Get smem layout for A, transformed A, and B + ( + self.smem_layout_a, + self.smem_layout_a_transform, + self.smem_layout_b, + ) = mixed_input_utils.compute_smem_layout( + tiled_mma, + self.mma_tiler, + self.a_dtype, + self.b_dtype, + self.num_load2trans_stage, + self.num_trans2mma_stage, + ) + # Get smem layout for scale tensor + self.smem_layout_scale_per_stage = None + self.smem_layout_scale = None + if cutlass.const_expr(self.scale_mode == TransformMode.ConvertScale): + # Get scale tile shape and smem layout for scale tensor + ( + self.scale_tile_shape, + self.smem_layout_scale_per_stage, + self.smem_layout_scale, + ) = mixed_input_utils.get_smem_layout_scale( + self.mma_tiler, + self.use_2cta_instrs, + self.scale_granularity_m, + self.scale_granularity_k, + self.scale_major_mode, + self.a_scale_dtype, + self.num_scale_load2trans_stage, + ) + + def _validate_inputs( + self, + a: cute.Tensor, + a_scale: Optional[cute.Tensor], + b: cute.Tensor, + c: cute.Tensor, + ) -> None: + """ + Validates input tensors and their properties. + + :param a: Input tensor A. + :type a: cute.Tensor + :param a_scale: Scale tensor for tensor A (None for ConvertOnly mode). + :type a_scale: Optional[cute.Tensor] + :param b: Input tensor B. + :type b: cute.Tensor + :param c: Output tensor C. + :type c: cute.Tensor + :raises ValueError: If inputs don't meet kernel requirements. + """ + # Validate scale tensor major mode + if cutlass.const_expr( + self.scale_mode == TransformMode.ConvertScale + and utils.LayoutEnum.from_tensor(a_scale).mma_major_mode() + != tcgen05.OperandMajorMode.MN + ): + raise ValueError("scale_major_mode must be M-major") + + @cute.jit + def __call__( + self, + a: cute.Tensor, + a_scale: Optional[cute.Tensor], # None for ConvertOnly mode + b: cute.Tensor, + c: cute.Tensor, + max_active_clusters: cutlass.Constexpr, + stream: cuda.CUstream, + ): + """ + Executes the Mixed Input GEMM operation. + + This method sets up the kernel parameters, computes the grid size, + defines the shared storage, and launches the kernel. + + The execution steps are as follows: + - Setup static attributes before smem/grid/tma computation. + - Setup TMA load/store atoms and tensors. + - Compute grid size with regard to hardware constraints. + - Define shared storage for kernel. + - Launch the kernel synchronously. + + :param a: Input tensor A. + :type a: cute.Tensor + :param a_scale: Scale tensor for tensor A (None for ConvertOnly mode). + :type a_scale: Optional[cute.Tensor] + :param b: Input tensor B. + :type b: cute.Tensor + :param c: Output tensor C. + :type c: cute.Tensor + :param max_active_clusters: Maximum number of active clusters to launch. + :type max_active_clusters: cutlass.Constexpr + :param stream: CUDA stream to launch the kernel on. + :type stream: cuda.CUstream + """ + self.a_dtype: type[cutlass.Numeric] = a.element_type + self.a_scale_dtype: type[cutlass.Numeric] = ( + a_scale.element_type + if self.scale_mode is TransformMode.ConvertScale + else None + ) + self.b_dtype: type[cutlass.Numeric] = b.element_type + self.c_dtype: type[cutlass.Numeric] = c.element_type + self.mma_dtype = self.b_dtype + + self.a_major_mode = utils.LayoutEnum.from_tensor(a).mma_major_mode() + self.scale_major_mode = ( + utils.LayoutEnum.from_tensor(a_scale).mma_major_mode() + if self.scale_mode is TransformMode.ConvertScale + else None + ) + self.b_major_mode = utils.LayoutEnum.from_tensor(b).mma_major_mode() + self.c_layout = utils.LayoutEnum.from_tensor(c) + if cutlass.const_expr(self.scale_mode == TransformMode.ConvertScale): + # Get gmem layout for scale tensor + self.gmem_layout_scale = mixed_input_utils.get_gmem_layout_scale( + a.shape, + self.scale_granularity_m, + self.scale_granularity_k, + self.scale_major_mode, + ) + + # Validate inputs + self._validate_inputs(a, a_scale, b, c) + + # Setup attributes that dependent on gemm inputs + self._setup_attributes() + + tiled_mma = sm100_utils.make_trivial_tiled_mma( + self.mma_dtype, + self.a_major_mode, + self.b_major_mode, + self.acc_dtype, + self.cta_group, + self.mma_tiler[:2], + self.transform_a_source, + ) + # Set up gmem copy atoms for A, scale, and B + a_op = mixed_input_utils.get_tma_atom_kind( + self.is_a_mcast, self.use_2cta_instrs, False + ) + b_op = mixed_input_utils.get_tma_atom_kind( + self.is_b_mcast, self.use_2cta_instrs, True + ) + a_scale_op = a_op + # Deduce TMA copy atom and TMA tensor for A, scale, and B + smem_layout_a_per_stage = cute.slice_(self.smem_layout_a, (None, None, None, 0)) + tma_atom_a, tma_tensor_a = cute.nvgpu.make_tiled_tma_atom_A( + a_op, + a, + smem_layout_a_per_stage, + self.mma_tiler, + tiled_mma, + self.cluster_layout_vmnk.shape, + internal_type=( + cutlass.TFloat32 if a.element_type is cutlass.Float32 else None + ), + ) + + tma_atom_scale, tma_tensor_scale = None, None + if cutlass.const_expr(self.scale_mode == TransformMode.ConvertScale): + # Partition smem layout for scale tensor to make it compatible with TMA atom + smem_layout_for_tma_atom = cute.get( + tiled_mma._thrfrg_A(self.smem_layout_scale_per_stage.outer), mode=[1] + ) + # ((MMA_M, MMA_K), REST_M, REST_K) + smem_layout_for_tma_atom = cute.dice( + smem_layout_for_tma_atom, + (1, (1,) * cute.rank(self.smem_layout_scale_per_stage.outer)), + ) + tma_atom_scale, tma_tensor_scale = cute.nvgpu.make_tiled_tma_atom_A( + a_scale_op, + cute.make_tensor(a_scale.iterator, self.gmem_layout_scale), + smem_layout_for_tma_atom, + # (SCALE_M, 1, SCALE_K) + (self.scale_tile_shape[0], 1, self.scale_tile_shape[1]), + tiled_mma, + self.cluster_layout_vmnk.shape, + internal_type=( + cutlass.TFloat32 + if a_scale.element_type is cutlass.Float32 + else None + ), + ) + + smem_layout_b_per_stage = cute.slice_(self.smem_layout_b, (None, None, None, 0)) + tma_atom_b, tma_tensor_b = cute.nvgpu.make_tiled_tma_atom_B( + b_op, + b, + smem_layout_b_per_stage, + self.mma_tiler, + tiled_mma, + self.cluster_layout_vmnk.shape, + internal_type=( + cutlass.TFloat32 if b.element_type is cutlass.Float32 else None + ), + ) + + # Calculate copy size for tensor A, B, and scale + a_copy_size = cute.size_in_bytes(self.a_dtype, smem_layout_a_per_stage) + b_copy_size = cute.size_in_bytes(self.b_dtype, smem_layout_b_per_stage) + a_scale_copy_size = ( + cute.size_in_bytes(self.a_scale_dtype, self.smem_layout_scale_per_stage) + if self.scale_mode is TransformMode.ConvertScale + else 0 + ) + + self.num_tma_load_bytes_a = a_copy_size + self.num_tma_load_bytes_b = b_copy_size * cute.size(tiled_mma.thr_id.shape) + self.num_tma_load_bytes_scale = a_scale_copy_size + self.tile_sched_params, grid = self._compute_grid( + c, + self.cta_tile_shape_mnk, + self.cluster_shape_mn, + max_active_clusters, + ) + + tma_atom_c = None + tma_tensor_c = None + c_smem_size = 0 + if cutlass.const_expr(self.use_tma_store): + epi_smem_layout = cute.slice_(self.c_smem_layout_staged, (None, None, 0)) + tma_atom_c, tma_tensor_c = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileS2GOp(), + c, + epi_smem_layout, + self.epi_tile, + ) + c_smem_size = cute.cosize(self.c_smem_layout_staged.outer) + + # Shared memory structure + a_smem_size = cute.cosize(self.smem_layout_a.outer) + b_smem_size = cute.cosize(self.smem_layout_b.outer) + a_transform_smem_size = ( + cute.cosize(self.smem_layout_a_transform.outer) + if self.transform_a_source == tcgen05.OperandSource.SMEM + else 0 + ) + a_scale_smem_size = ( + cute.cosize(self.smem_layout_scale.outer) + if self.scale_mode is TransformMode.ConvertScale + else 0 + ) + + @cute.struct + class SharedStorage: + a_load2trans_full_mbar_ptr: cute.struct.MemRange[ + cutlass.Int64, self.num_load2trans_stage + ] + a_load2trans_empty_mbar_ptr: cute.struct.MemRange[ + cutlass.Int64, self.num_load2trans_stage + ] + a_scale_load2trans_full_mbar_ptr: cute.struct.MemRange[ + cutlass.Int64, self.num_scale_load2trans_stage + ] + a_scale_load2trans_empty_mbar_ptr: cute.struct.MemRange[ + cutlass.Int64, self.num_scale_load2trans_stage + ] + a_trans2mma_full_mbar_ptr: cute.struct.MemRange[ + cutlass.Int64, self.num_trans2mma_stage + ] + a_trans2mma_empty_mbar_ptr: cute.struct.MemRange[ + cutlass.Int64, self.num_trans2mma_stage + ] + b_load2mma_full_mbar_ptr: cute.struct.MemRange[ + cutlass.Int64, self.num_load2trans_stage + ] + b_load2mma_empty_mbar_ptr: cute.struct.MemRange[ + cutlass.Int64, self.num_load2trans_stage + ] + acc_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage] + acc_empty_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage] + tmem_dealloc_mbar_ptr: cutlass.Int64 + tmem_holding_buf: cutlass.Int32 + # Tensor buffers + # (EPI_TILE_M, EPI_TILE_N, STAGE) + smem_C: cute.struct.Align[ + cute.struct.MemRange[self.c_dtype, c_smem_size], + self.smem_buffer_align_bytes, + ] + # (MMA, MMA_M, MMA_K, STAGE) + smem_A: cute.struct.Align[ + cute.struct.MemRange[self.a_dtype, a_smem_size], + self.smem_buffer_align_bytes, + ] + # (MMA, MMA_N, MMA_K, STAGE) + smem_B: cute.struct.Align[ + cute.struct.MemRange[self.b_dtype, b_smem_size], + self.smem_buffer_align_bytes, + ] + # (MMA, MMA_M, MMA_K, STAGE) + smem_A_transform: cute.struct.Align[ + cute.struct.MemRange[self.mma_dtype, a_transform_smem_size], + self.smem_buffer_align_bytes, + ] + # (MMA, MMA_M_SCALE, MMA_K_SCALE, STAGE) + smem_A_scale: cute.struct.Align[ + cute.struct.MemRange[self.mma_dtype, a_scale_smem_size], + self.smem_buffer_align_bytes, + ] + + self.shared_storage = SharedStorage + + # Launch kernel + self.kernel( + tiled_mma, + tma_atom_a, + tma_tensor_a, + tma_atom_scale, + tma_tensor_scale, + tma_atom_b, + tma_tensor_b, + tma_atom_c, + tma_tensor_c if self.use_tma_store else c, + self.cluster_layout_vmnk, + self.smem_layout_a, + self.smem_layout_scale, + self.smem_layout_a_transform, + self.smem_layout_b, + self.c_smem_layout_staged, + self.epi_tile, + self.tile_sched_params, + ).launch( + grid=grid, + block=[self.threads_per_cta, 1, 1], + cluster=(*self.cluster_shape_mn, 1), + min_blocks_per_mp=1, + stream=stream, + ) + return + + # GPU device kernel + @cute.kernel + def kernel( + self, + tiled_mma: cute.TiledMma, + tma_atom_a: cute.CopyAtom, + mA_mkl: cute.Tensor, + tma_atom_s: Optional[cute.CopyAtom], + mS_mkl: Optional[cute.Tensor], + tma_atom_b: cute.CopyAtom, + mB_nkl: cute.Tensor, + tma_atom_c: Optional[cute.CopyAtom], + mC_mnl: cute.Tensor, + cluster_layout_vmnk: cute.Layout, + a_smem_layout: cute.ComposedLayout, + scale_smem_layout: cute.ComposedLayout, + a_smem_layout_transform: cute.ComposedLayout, + b_smem_layout: cute.ComposedLayout, + c_smem_layout_staged: cute.ComposedLayout, + epi_tile: cute.Tile, + tile_sched_params: utils.PersistentTileSchedulerParams, + ): + """ + GPU device kernel performing the Persistent Mixed-Input GEMM computation. + """ + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + tidx, _, _ = cute.arch.thread_idx() + bidx, bidy, bidz = cute.arch.block_idx() + # Prefetch TMA descriptors + if warp_idx == self.epilog_warp_id[0]: + cpasync.prefetch_descriptor(tma_atom_a) + cpasync.prefetch_descriptor(tma_atom_b) + if cutlass.const_expr(self.scale_mode == TransformMode.ConvertScale): + cpasync.prefetch_descriptor(tma_atom_s) + if cutlass.const_expr(self.use_tma_store): + cpasync.prefetch_descriptor(tma_atom_c) + + use_2cta_instrs = cute.size(tiled_mma.thr_id.shape) == 2 + bidx, bidy, bidz = cute.arch.block_idx() + # Compute how many k_tiles share the same scale + num_k_tiles_per_scale = self.scale_granularity_k // self.cta_tile_shape_mnk[2] + + mma_tile_coord_v = bidx % cute.size(tiled_mma.thr_id.shape) + is_leader_cta = mma_tile_coord_v == 0 + cta_rank_in_cluster = cute.arch.make_warp_uniform( + cute.arch.block_idx_in_cluster() + ) + block_in_cluster_coord_vmnk = cluster_layout_vmnk.get_flat_coord( + cta_rank_in_cluster + ) + tidx, _, _ = cute.arch.thread_idx() + + smem = utils.SmemAllocator() + storage = smem.allocate(self.shared_storage) + + # Initialize load2transform pipeline, which tracks the dependencies between TMA's loading + # of A and B, and the transformation of A and MMA's consumption + transform_thread_idx = ( + tidx - 32 * self.transform_warp_id[0] + if tidx >= 32 * self.transform_warp_id[0] + else tidx + ) + a_load2trans_pipeline = pipeline.PipelineTmaAsync.create( + barrier_storage=storage.a_load2trans_full_mbar_ptr.data_ptr(), + num_stages=self.num_load2trans_stage, + producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), + consumer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, + self.num_mcast_ctas_a * len(self.transform_warp_id), + ), + tx_count=self.num_tma_load_bytes_a, + cta_layout_vmnk=cluster_layout_vmnk, + tidx=transform_thread_idx, + mcast_mode_mn=(1, 0), # multicast for A will only happen on the M-mode + defer_sync=True, + ) + # Initialize scale_load2trans pipeline, which tracks the dependencies between TMA's loading + # of scale, and the transformation of A + scale_load2trans_pipeline = None + if cutlass.const_expr(self.scale_mode == TransformMode.ConvertScale): + num_producers_a_scale = self.num_mcast_ctas_a + scale_load2trans_pipeline = pipeline.PipelineTmaAsync.create( + barrier_storage=storage.a_scale_load2trans_full_mbar_ptr.data_ptr(), + num_stages=self.num_scale_load2trans_stage, + producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), + consumer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, + num_producers_a_scale + * len(self.transform_warp_id) + * num_k_tiles_per_scale, + ), + tx_count=self.num_tma_load_bytes_scale, + cta_layout_vmnk=cluster_layout_vmnk, + tidx=transform_thread_idx, + mcast_mode_mn=( + 1, + 0, + ), # multicast for scale_a will only happen on the M-mode + defer_sync=True, + ) + # Initialize transform2mma pipeline, which tracks the dependencies between the transformation + # of A and MMA's consumption of transformed A + cta_v_size = cute.size(cluster_layout_vmnk, mode=[0]) + trans2mma_pipeline = pipeline.PipelineAsyncUmma.create( + barrier_storage=storage.a_trans2mma_full_mbar_ptr.data_ptr(), + num_stages=self.num_trans2mma_stage, + producer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, + 32 * len(self.transform_warp_id) * cta_v_size, + ), + consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ) + # Initialize pipeline for tensor B load to MMA + # MMA warp informs TMA warp to proceed to load next tile of B tensor + b_load2mma_pipeline = pipeline.PipelineTmaUmma.create( + barrier_storage=storage.b_load2mma_full_mbar_ptr.data_ptr(), + num_stages=self.num_load2trans_stage, + producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), + consumer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, self.num_mcast_ctas_b + ), + tx_count=self.num_tma_load_bytes_b, + cta_layout_vmnk=cluster_layout_vmnk, + mcast_mode_mn=(0, 1), # multicast for B will only happen on the N-mode + defer_sync=True, + ) + # Initialize accumulator pipeline, which tracks the dependencies between + # MMA's computation of accumulators and epilogue warps' consumption of accumulators + acc_pipeline = pipeline.PipelineUmmaAsync.create( + barrier_storage=storage.acc_full_mbar_ptr.data_ptr(), + num_stages=self.num_acc_stage, + producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), + consumer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, cta_v_size * len(self.epilog_warp_id) + ), + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ) + + # Tensor memory dealloc barrier init + tmem = utils.TmemAllocator( + storage.tmem_holding_buf, + barrier_for_retrieve=self.tmem_ptr_sync_barrier, + allocator_warp_id=self.epilog_warp_id[0], + is_two_cta=use_2cta_instrs, + two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar_ptr, + ) + + # Cluster arrive after barrier init + pipeline_init_arrive(cluster_shape_mn=self.cluster_shape_mn, is_relaxed=True) + + # Setup smem tensor A/scale/B/C + sC = ( + storage.smem_C.get_tensor( + c_smem_layout_staged.outer, swizzle=c_smem_layout_staged.inner + ) + if self.use_tma_store + else None + ) + sA_input = storage.smem_A.get_tensor( + a_smem_layout.outer, swizzle=a_smem_layout.inner + ) + sS_input = ( + storage.smem_A_scale.get_tensor( + scale_smem_layout.outer, swizzle=scale_smem_layout.inner + ) + if self.scale_mode is TransformMode.ConvertScale + else None + ) + sB_input = storage.smem_B.get_tensor( + b_smem_layout.outer, swizzle=b_smem_layout.inner + ) + sA_transform = None + # Get smem tensor for transformed A when transform_a_source is SMEM + if cutlass.const_expr(self.transform_a_source == tcgen05.OperandSource.SMEM): + sA_transform = storage.smem_A_transform.get_tensor( + a_smem_layout_transform.outer, swizzle=a_smem_layout_transform.inner + ) + + # Compute multicast mask for A/B buffer full + a_full_mcast_mask = None + b_full_mcast_mask = None + s_full_mcast_mask = None + if cutlass.const_expr(self.is_a_mcast or self.is_b_mcast or use_2cta_instrs): + a_full_mcast_mask = cpasync.create_tma_multicast_mask( + cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=2 + ) + # scale tensor share the same multicast mask with A tensor + s_full_mcast_mask = a_full_mcast_mask + b_full_mcast_mask = cpasync.create_tma_multicast_mask( + cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=1 + ) + + # local_tile partition global tensors + # (bM, bK, loopM, loopK, loopL) + gA_mkl = cute.local_tile( + mA_mkl, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None) + ) + # (bM, bK, loopM, loopK, loopL) + gS_mkl = ( + cute.local_tile( + mS_mkl, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None) + ) + if self.scale_mode is TransformMode.ConvertScale + else None + ) + # (bN, bK, loopN, loopK, loopL) + gB_nkl = cute.local_tile( + mB_nkl, cute.slice_(self.mma_tiler, (0, None, None)), (None, None, None) + ) + # (bM, bN, loopM, loopN, loopL) + gC_mnl = cute.local_tile( + mC_mnl, cute.slice_(self.mma_tiler, (None, None, 0)), (None, None, None) + ) + k_tile_cnt = cute.size(gA_mkl, mode=[3]) + + # Partition global tensor for TiledMMA_A/B/C + thr_mma = tiled_mma.get_slice(mma_tile_coord_v) + # (MMA, MMA_M, MMA_K, loopM, loopK, loopL) + tCgA = thr_mma.partition_A(gA_mkl) + # (MMA, MMA_M, MMA_K, loopM, loopK, loopL) + tCgS = ( + thr_mma.partition_A(gS_mkl) + if self.scale_mode is TransformMode.ConvertScale + else None + ) + # (MMA, MMA_N, MMA_K, loopN, loopK, loopL) + tCgB = thr_mma.partition_B(gB_nkl) + # (MMA, MMA_M, MMA_N, loopM, loopN, loopL) + tCgC = thr_mma.partition_C(gC_mnl) + + # Setup copy atom to load A from shared memory for further transformation + copy_atom_a_input = ( + cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), self.a_dtype, num_bits_per_copy=32 + ) + if self.scale_mode is TransformMode.ConvertScale + else None + ) + a_smem_shape = tiled_mma.partition_shape_A( + cute.dice(self.mma_tiler, (1, None, 1)) + ) + # Setup copy atom to store transformed A into tensor memory or shared memory + copy_atom_a_transform = mixed_input_utils.get_copy_atom_a_transform( + self.mma_dtype, + self.use_2cta_instrs, + self.transform_a_source, + a_smem_shape, + self.a_dtype, + ) + + # Partition global/shared tensor for TMA load A/B + # TMA load A partition_S/D + a_cta_layout = cute.make_layout( + cute.slice_(cluster_layout_vmnk, (0, 0, None, 0)).shape + ) + # ((atom_v, rest_v), STAGE) + # ((atom_v, rest_v), loopM, loopK, loopL) + tAsA, tAgA = cpasync.tma_partition( + tma_atom_a, + block_in_cluster_coord_vmnk[2], + a_cta_layout, + cute.group_modes(sA_input, 0, 3), + cute.group_modes(tCgA, 0, 3), + ) + + tCsS = None + tSsS = None + tSgS = None + if cutlass.const_expr(self.scale_mode == TransformMode.ConvertScale): + thr_mma_leader_cta = tiled_mma.get_slice(0) + # (MMA, MMA_M, MMA_K, STAGE) + tCsS = thr_mma_leader_cta.partition_A(sS_input) + # ((atom_v, rest_v), STAGE) + # ((atom_v, rest_v), loopM, loopK, loopL) + tSsS, tSgS = mixed_input_utils.scale_tma_partition( + tCsS, + tCgS, + tma_atom_s, + block_in_cluster_coord_vmnk, + a_cta_layout, + ) + + # TMA load B partition_S/D + b_cta_layout = cute.make_layout( + cute.slice_(cluster_layout_vmnk, (0, None, 0, 0)).shape + ) + # ((atom_v, rest_v), STAGE) + # ((atom_v, rest_v), loopM, loopK, loopL) + tBsB, tBgB = cpasync.tma_partition( + tma_atom_b, + block_in_cluster_coord_vmnk[1], + b_cta_layout, + cute.group_modes(sB_input, 0, 3), + cute.group_modes(tCgB, 0, 3), + ) + + # (MMA, MMA_N, MMA_K, STAGE) + tCrB = tiled_mma.make_fragment_B(sB_input) + # (MMA, MMA_M, MMA_N) + acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2]) + tCtAcc_fake = tiled_mma.make_fragment_C( + cute.append(acc_shape, self.num_acc_stage) + ) + + # Cluster wait before TMEM alloc and ensure pipelines are ready + pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mn) + + # TMEM allocation + tmem.allocate(self.num_tmem_alloc_cols) + tmem.wait_for_alloc() + # Get the pointer to the TMEM buffer + tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + accumulators = cute.make_tensor(tmem_ptr, tCtAcc_fake.layout) + + tCrA = None + if cutlass.const_expr(self.transform_a_source == tcgen05.OperandSource.TMEM): + tmem_ptr_transform = cute.recast_ptr( + accumulators.iterator + self.num_acc_tmem_cols, dtype=self.mma_dtype + ) + tCrA = cute.make_tensor( + tmem_ptr_transform, + tiled_mma.make_fragment_A(a_smem_layout_transform.outer).layout, + ) + else: + tCrA = tiled_mma.make_fragment_A(sA_transform) + + # Specialized TMA load warp for A/B tensor + if warp_idx == self.tma_warp_id: + cute.arch.setmaxregister_decrease(self.num_regs_tma_warps) + # Persistent tile scheduling loop + tile_sched = utils.StaticPersistentTileScheduler.create( + tile_sched_params, (bidx, bidy, bidz), cute.arch.grid_dim() + ) + work_tile = tile_sched.initial_work_tile_info() + a_load2trans_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.num_load2trans_stage + ) + b_load2mma_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.num_load2trans_stage + ) + while work_tile.is_valid_tile: + # Get tile coord from tile scheduler + cur_tile_coord = work_tile.tile_idx + mma_tile_coord_mnl = ( + cur_tile_coord[0] // cute.size(tiled_mma.thr_id.shape), + cur_tile_coord[1], + cur_tile_coord[2], + ) + tAgA_slice = tAgA[ + (None, mma_tile_coord_mnl[0], None, mma_tile_coord_mnl[2]) + ] + tBgB_slice = tBgB[ + (None, mma_tile_coord_mnl[1], None, mma_tile_coord_mnl[2]) + ] + + a_load2trans_producer_state.reset_count() + peek_load2trans_empty_status = cutlass.Boolean(1) + if a_load2trans_producer_state.count < k_tile_cnt: + peek_load2trans_empty_status = ( + a_load2trans_pipeline.producer_try_acquire( + a_load2trans_producer_state + ) + ) + b_load2mma_producer_state.reset_count() + for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1): + a_load2trans_pipeline.producer_acquire( + a_load2trans_producer_state, peek_load2trans_empty_status + ) + b_load2mma_pipeline.producer_acquire(b_load2mma_producer_state) + # TMA load A/B + cute.copy( + tma_atom_a, + tAgA_slice[(None, a_load2trans_producer_state.count)], + tAsA[(None, a_load2trans_producer_state.index)], + tma_bar_ptr=a_load2trans_pipeline.producer_get_barrier( + a_load2trans_producer_state + ), + mcast_mask=a_full_mcast_mask, + ) + cute.copy( + tma_atom_b, + tBgB_slice[(None, b_load2mma_producer_state.count)], + tBsB[(None, b_load2mma_producer_state.index)], + tma_bar_ptr=b_load2mma_pipeline.producer_get_barrier( + b_load2mma_producer_state + ), + mcast_mask=b_full_mcast_mask, + ) + a_load2trans_pipeline.producer_commit(a_load2trans_producer_state) + b_load2mma_pipeline.producer_commit(b_load2mma_producer_state) + a_load2trans_producer_state.advance() + b_load2mma_producer_state.advance() + if a_load2trans_producer_state.count < k_tile_cnt: + peek_load2trans_empty_status = ( + a_load2trans_pipeline.producer_try_acquire( + a_load2trans_producer_state + ) + ) + # Advance to next tile + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + # Wait A/B buffer empty + a_load2trans_pipeline.producer_tail(a_load2trans_producer_state) + b_load2mma_pipeline.producer_tail(b_load2mma_producer_state) + + # Specialized TMA load for scale tensor + if warp_idx == self.scale_tma_warp_id: + cute.arch.setmaxregister_decrease(self.num_regs_tma_warps) + if cutlass.const_expr(self.scale_mode == TransformMode.ConvertScale): + # Persistent tile scheduling loop + tile_sched = utils.StaticPersistentTileScheduler.create( + tile_sched_params, (bidx, bidy, bidz), cute.arch.grid_dim() + ) + work_tile = tile_sched.initial_work_tile_info() + scale_load2trans_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.num_scale_load2trans_stage + ) + scale_k_tile_cnt = cute.size(mS_mkl.layout.shape[1][1]) + + while work_tile.is_valid_tile: + cur_tile_coord = work_tile.tile_idx + mma_tile_coord_mnl = ( + cur_tile_coord[0] // cute.size(tiled_mma.thr_id.shape), + cur_tile_coord[1], + cur_tile_coord[2], + ) + # ((atom_v, rest_v), RestK) + tSgS_slice = tSgS[ + (None, mma_tile_coord_mnl[0], None, mma_tile_coord_mnl[2]) + ] + # Filter zeros in rest mode + rest_filtered = cute.filter_zeros(tSgS_slice[(0, None)].layout) + tSgS_slice_filtered = cute.make_tensor( + tSgS_slice.iterator, + cute.make_layout( + (tSgS_slice.layout[0].shape, rest_filtered.shape), + stride=(tSgS_slice.layout[0].stride, rest_filtered.stride), + ), + ) + + scale_load2trans_producer_state.reset_count() + peek_scale_load2trans_empty_status = cutlass.Boolean(1) + if scale_load2trans_producer_state.count < scale_k_tile_cnt: + peek_scale_load2trans_empty_status = ( + scale_load2trans_pipeline.producer_try_acquire( + scale_load2trans_producer_state + ) + ) + for k_tile in cutlass.range(0, scale_k_tile_cnt, 1, unroll=1): + scale_load2trans_pipeline.producer_acquire( + scale_load2trans_producer_state, + peek_scale_load2trans_empty_status, + ) + # TMA load scale + cute.copy( + tma_atom_s, + tSgS_slice_filtered[ + (None, scale_load2trans_producer_state.count) + ], + tSsS[(None, scale_load2trans_producer_state.index)], + tma_bar_ptr=scale_load2trans_pipeline.producer_get_barrier( + scale_load2trans_producer_state + ), + mcast_mask=s_full_mcast_mask, + ) + + scale_load2trans_producer_state.advance() + peek_scale_load2trans_empty_status = cutlass.Boolean(1) + if scale_load2trans_producer_state.count < scale_k_tile_cnt: + peek_scale_load2trans_empty_status = ( + scale_load2trans_pipeline.producer_try_acquire( + scale_load2trans_producer_state + ) + ) + # Advance to next tile + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + # Wait scale buffer empty + scale_load2trans_pipeline.producer_tail(scale_load2trans_producer_state) + + # Specialized transform warps + if warp_idx >= self.transform_warp_id[0]: + cute.arch.setmaxregister_increase(self.num_regs_transform_warps) + transform_local_tidx = tidx - 32 * self.transform_warp_id[0] + # Partition tensors for transform input and output and set up the copy atom + # used for loading and storing transformed A tensor + ( + src_copy_a, + dst_copy_a, + tAsA_input, + tAsA_transform, + ) = mixed_input_utils.transform_partition( + self.transform_a_source, + self.scale_mode, + copy_atom_a_input, + copy_atom_a_transform, + sA_input, + ( + tCrA + if self.transform_a_source == tcgen05.OperandSource.TMEM + else sA_transform + ), + transform_local_tidx, + ) + # make rmem tensor for input A and transformed A + tArA = cute.make_rmem_tensor( + tAsA_input[(None, None, None, None, 0)].shape, tAsA_input.element_type + ) + tArA_transform = cute.make_rmem_tensor( + tAsA_input[(None, None, None, None, 0)].shape, self.mma_dtype + ) + # Partition scale tensor + smem_thr_copy_S = None + tSsS_trans = None + tSrS_copy = None + tSrS = None + if cutlass.const_expr(self.scale_mode == TransformMode.ConvertScale): + smem_thr_copy_S, tSsS_trans, tSrS_copy, tSrS = ( + mixed_input_utils.scale_partition( + src_copy_a, tCsS, transform_local_tidx, self.mma_dtype + ) + ) + assert cute.size(tSrS, mode=[0]) == cute.size(tArA, mode=[0]), ( + "tSrS and tArA have different leading dimension" + ) + assert cute.size(tSrS) == cute.size(tArA), ( + "tSrS and tArA have different shape" + ) + # Deduce a sub-tile size and tile tensors + transform_tiler_size = min( + cute.size(cute.coalesce(tAsA_input.layout), mode=[0]), 64 + ) + transform_tiler = cute.make_layout(transform_tiler_size) + tArA_load = cute.flat_divide(tArA, transform_tiler) + tArA_load = cute.group_modes(tArA_load, 1, cute.rank(tArA_load)) + tSrS_load = ( + cute.flat_divide(tSrS, transform_tiler) + if self.scale_mode is TransformMode.ConvertScale + else None + ) + tSrS_load = ( + cute.group_modes(tSrS_load, 1, cute.rank(tSrS_load)) + if self.scale_mode is TransformMode.ConvertScale + else None + ) + tArA_transform_store = cute.flat_divide(tArA_transform, transform_tiler) + tArA_transform_store = cute.group_modes( + tArA_transform_store, 1, cute.rank(tArA_transform_store) + ) + + tile_sched = utils.StaticPersistentTileScheduler.create( + tile_sched_params, (bidx, bidy, bidz), cute.arch.grid_dim() + ) + work_tile = tile_sched.initial_work_tile_info() + a_load2trans_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, + self.num_load2trans_stage, + ) + scale_load2trans_consumer_state = ( + pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, + self.num_scale_load2trans_stage, + ) + if self.scale_mode is TransformMode.ConvertScale + else None + ) + trans2mma_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, + self.num_trans2mma_stage, + ) + while work_tile.is_valid_tile: + a_load2trans_consumer_state.reset_count() + peek_load2trans_full_status = cutlass.Boolean(1) + if a_load2trans_consumer_state.count < k_tile_cnt: + peek_load2trans_full_status = ( + a_load2trans_pipeline.consumer_try_wait( + a_load2trans_consumer_state + ) + ) + peek_scale_load2trans_full_status = cutlass.Boolean(1) + if cutlass.const_expr(self.scale_mode == TransformMode.ConvertScale): + scale_load2trans_consumer_state.reset_count() + peek_scale_load2trans_full_status = ( + scale_load2trans_pipeline.consumer_try_wait( + scale_load2trans_consumer_state + ) + ) + trans2mma_producer_state.reset_count() + peek_trans2mma_empty_status = cutlass.Boolean(1) + if trans2mma_producer_state.count < k_tile_cnt: + peek_trans2mma_empty_status = ( + trans2mma_pipeline.producer_try_acquire( + trans2mma_producer_state + ) + ) + + for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1): + a_load2trans_pipeline.consumer_wait( + a_load2trans_consumer_state, peek_load2trans_full_status + ) + # Load A from shared memory + tAsA_input_slice = tAsA_input[ + (None, None, None, None, a_load2trans_consumer_state.index) + ] + tAsA_input_slice = cute.flat_divide( + tAsA_input_slice, transform_tiler + ) + tAsA_input_slice = cute.group_modes( + tAsA_input_slice, 1, cute.rank(tAsA_input_slice) + ) + if cutlass.const_expr( + self.scale_mode == TransformMode.ConvertScale + ): + scale_load2trans_pipeline.consumer_wait( + scale_load2trans_consumer_state, + peek_scale_load2trans_full_status, + ) + trans2mma_pipeline.producer_acquire( + trans2mma_producer_state, peek_trans2mma_empty_status + ) + # load scale tensor when needed + if cutlass.const_expr( + self.scale_mode == TransformMode.ConvertScale + ): + if k_tile % num_k_tiles_per_scale == 0: + tSsS_slice = tSsS_trans[ + ( + None, + None, + None, + None, + scale_load2trans_consumer_state.index, + ) + ] + tSsS_slice_filtered = cute.make_tensor( + tSsS_slice.iterator, + cute.filter_zeros(tSsS_slice.layout), + ) + cute.autovec_copy(tSsS_slice_filtered, tSrS_copy) + cur_scale_load2trans_consumer_state = ( + scale_load2trans_consumer_state.clone() + ) + if (k_tile + 1) % num_k_tiles_per_scale == 0: + scale_load2trans_consumer_state.advance() + + cur_a_load2trans_consumer_state = ( + a_load2trans_consumer_state.clone() + ) + for idx in cutlass.range_constexpr(cute.size(tArA_load, mode=[1])): + # Load A from shared memory + cute.autovec_copy( + tAsA_input_slice[(None, idx)], + tArA_load[(None, idx)], + ) + if cutlass.const_expr( + idx == cute.size(tArA_load, mode=[1]) - 1 + ): + a_load2trans_consumer_state.advance() + if a_load2trans_consumer_state.count < k_tile_cnt: + peek_load2trans_full_status = ( + a_load2trans_pipeline.consumer_try_wait( + a_load2trans_consumer_state + ) + ) + if cutlass.const_expr( + self.scale_mode == TransformMode.ConvertScale + ): + peek_scale_load2trans_full_status = ( + scale_load2trans_pipeline.consumer_try_wait( + scale_load2trans_consumer_state + ) + ) + # Convert it to mma dtype + tensor_transformed = mixed_input_utils.cvt_tensor_a( + tArA_load[(None, idx)], self.mma_dtype, self.shuffle_a + ) + if cutlass.const_expr( + self.scale_mode == TransformMode.ConvertScale + ): + scale = cute.TensorSSA( + tSrS_load[(None, idx)].load(), + tensor_transformed.shape, + self.mma_dtype, + ) + # Apply scale + tensor_transformed = tensor_transformed * scale + tArA_transform_store[(None, idx)].store(tensor_transformed) + # Store transformed A to tensor memory or shared memory + mixed_input_utils.store_transformed_a( + tArA_transform, + tAsA_transform[ + (None, None, None, None, trans2mma_producer_state.index) + ], + dst_copy_a, + ) + # Ensure all transform threads have finished the copy and reached the fence + self.transform_sync_barrier.arrive_and_wait() + if cutlass.const_expr( + self.transform_a_source == tcgen05.OperandSource.TMEM + ): + cute.arch.fence_view_async_tmem_store() + else: + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) + # Signal the completion of transformation + if cutlass.const_expr( + self.scale_mode == TransformMode.ConvertScale + ): + scale_load2trans_pipeline.consumer_release( + cur_scale_load2trans_consumer_state + ) + a_load2trans_pipeline.consumer_release( + cur_a_load2trans_consumer_state + ) + # Signal the completion of transformation + trans2mma_pipeline.producer_commit(trans2mma_producer_state) + trans2mma_producer_state.advance() + if trans2mma_producer_state.count < k_tile_cnt: + peek_trans2mma_empty_status = ( + trans2mma_pipeline.producer_try_acquire( + trans2mma_producer_state + ) + ) + # Advance to next tile + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + # Wait a_transform buffer empty + trans2mma_pipeline.producer_tail(trans2mma_producer_state) + + # Specialized MMA warp + if warp_idx == self.mma_warp_id: + cute.arch.setmaxregister_decrease(self.num_regs_mma_warp) + tCtAcc_base = accumulators + # Persistent tile scheduling loop + tile_sched = utils.StaticPersistentTileScheduler.create( + tile_sched_params, (bidx, bidy, bidz), cute.arch.grid_dim() + ) + work_tile = tile_sched.initial_work_tile_info() + trans2mma_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_trans2mma_stage + ) + b_load2mma_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_load2trans_stage + ) + acc_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.num_acc_stage + ) + while work_tile.is_valid_tile: + cur_tile_coord = work_tile.tile_idx + # (MMA, MMA_M, MMA_N) + tCtAcc = tCtAcc_base[(None, None, None, acc_producer_state.index)] + b_load2mma_consumer_state.reset_count() + trans2mma_consumer_state.reset_count() + peek_trans2mma_full_status = cutlass.Boolean(1) + if is_leader_cta: + if trans2mma_consumer_state.count < k_tile_cnt: + peek_trans2mma_full_status = ( + trans2mma_pipeline.consumer_try_wait( + trans2mma_consumer_state + ) + ) + acc_pipeline.producer_acquire(acc_producer_state) + + tiled_mma.set(tcgen05.Field.ACCUMULATE, False) + # Mma mainloop + for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1): + trans2mma_pipeline.consumer_wait( + trans2mma_consumer_state, peek_trans2mma_full_status + ) + b_load2mma_pipeline.consumer_wait(b_load2mma_consumer_state) + num_kblocks = cute.size(tCrA, mode=[2]) + for kblock_idx in cutlass.range(num_kblocks, unroll_full=True): + kblock_coord_a = ( + None, + None, + kblock_idx, + trans2mma_consumer_state.index, + ) + kblock_coord_b = ( + None, + None, + kblock_idx, + b_load2mma_consumer_state.index, + ) + + cute.gemm( + tiled_mma, + tCtAcc, + tCrA[kblock_coord_a], + tCrB[kblock_coord_b], + tCtAcc, + ) + # Enable accumulate on tCtAcc after first kblock + tiled_mma.set(tcgen05.Field.ACCUMULATE, True) + trans2mma_pipeline.consumer_release(trans2mma_consumer_state) + b_load2mma_pipeline.consumer_release(b_load2mma_consumer_state) + trans2mma_consumer_state.advance() + b_load2mma_consumer_state.advance() + peek_trans2mma_full_status = cutlass.Boolean(1) + if trans2mma_consumer_state.count < k_tile_cnt: + peek_trans2mma_full_status = ( + trans2mma_pipeline.consumer_try_wait( + trans2mma_consumer_state + ) + ) + # Async arrive accumulator buffer full + acc_pipeline.producer_commit(acc_producer_state) + acc_producer_state.advance() + + # Advance to next tile + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + # Wait for accumulator buffer empty + acc_pipeline.producer_tail(acc_producer_state) + + # Specialized epilogue warps + if warp_idx < self.mma_warp_id: + cute.arch.setmaxregister_increase(self.num_regs_epilogue_warps) + epi_tidx = tidx + tCtAcc_base = accumulators + # Partition for epilogue + ( + tiled_copy_t2r, + tTR_tAcc_base, + tTR_rAcc, + ) = mixed_input_utils.epilog_tmem_copy_and_partition( + self.cta_tile_shape_mnk, + self.c_layout, + self.c_dtype, + self.acc_dtype, + epi_tidx, + tCtAcc_base, + tCgC, + epi_tile, + self.use_2cta_instrs, + ) + + tTR_rC = None + tiled_copy_r2s = None + simt_atom = None + tRS_rC = None + tRS_sC = None + bSG_sC = None + bSG_gC_partitioned = None + tTR_gC_partitioned = None + if cutlass.const_expr(self.use_tma_store): + tTR_rC = cute.make_rmem_tensor(tTR_rAcc.shape, self.c_dtype) + tiled_copy_r2s, tRS_rC, tRS_sC = ( + mixed_input_utils.epilog_smem_copy_and_partition( + self.c_layout, + self.c_dtype, + self.acc_dtype, + tiled_copy_t2r, + tTR_rC, + epi_tidx, + sC, + ) + ) + ( + tma_atom_c, + bSG_sC, + bSG_gC_partitioned, + ) = self.epilog_gmem_copy_and_partition( + epi_tidx, tma_atom_c, tCgC, epi_tile, sC + ) + else: + ( + simt_atom, + tTR_rC, + tTR_gC_partitioned, + ) = self.epilog_gmem_copy_and_partition( + epi_tidx, tiled_copy_t2r, tCgC, epi_tile, sC + ) + # Persistent tile scheduling loop + tile_sched = utils.StaticPersistentTileScheduler.create( + tile_sched_params, (bidx, bidy, bidz), cute.arch.grid_dim() + ) + work_tile = tile_sched.initial_work_tile_info() + acc_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_acc_stage + ) + + c_pipeline = None + if cutlass.const_expr(self.use_tma_store): + c_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + 32 * len(self.epilog_warp_id), + ) + c_pipeline = pipeline.PipelineTmaStore.create( + num_stages=self.num_c_stage, + producer_group=c_producer_group, + ) + + while work_tile.is_valid_tile: + cur_tile_coord = work_tile.tile_idx + mma_tile_coord_mnl = ( + cur_tile_coord[0] // cute.size(tiled_mma.thr_id.shape), + cur_tile_coord[1], + cur_tile_coord[2], + ) + + bSG_gC = None + tTR_gC = None + if cutlass.const_expr(self.use_tma_store): + bSG_gC = bSG_gC_partitioned[(None, None, None, *mma_tile_coord_mnl)] + else: + tTR_gC = tTR_gC_partitioned[ + (None, None, None, None, None, *mma_tile_coord_mnl) + ] + + tTR_tAcc = tTR_tAcc_base[ + (None, None, None, None, None, acc_consumer_state.index) + ] + # Wait for accumulator buffer full + acc_pipeline.consumer_wait(acc_consumer_state) + + tTR_tAcc = cute.group_modes(tTR_tAcc, 3, cute.rank(tTR_tAcc)) + if cutlass.const_expr(self.use_tma_store): + bSG_gC = cute.group_modes(bSG_gC, 1, cute.rank(bSG_gC)) + else: + tTR_gC = cute.group_modes(tTR_gC, 3, cute.rank(tTR_gC)) + + # Store accumulator to global memory in subtiles + subtile_cnt = cute.size(tTR_tAcc.shape, mode=[3]) + num_prev_subtiles = tile_sched.num_tiles_executed * subtile_cnt + for subtile_idx in cutlass.range(subtile_cnt): + # Load accumulator from tensor memory buffer to register + tTR_tAcc_mn = tTR_tAcc[(None, None, None, subtile_idx)] + cute.copy(tiled_copy_t2r, tTR_tAcc_mn, tTR_rAcc) + if cutlass.const_expr(self.use_tma_store): + # Convert to C type + acc_vec = tiled_copy_r2s.retile(tTR_rAcc).load() + acc_vec = acc_vec.to(self.c_dtype) + tRS_rC.store(acc_vec) + c_buffer = (num_prev_subtiles + subtile_idx) % self.num_c_stage + # Store C to shared memory + cute.copy( + tiled_copy_r2s, + tRS_rC, + tRS_sC[(None, None, None, c_buffer)], + ) + # Fence and barrier to make sure shared memory store is visible to TMA store + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) + self.epilog_sync_barrier.arrive_and_wait() + # TMA store C to global memory + if warp_idx == self.epilog_warp_id[0]: + cute.copy( + tma_atom_c, + bSG_sC[(None, c_buffer)], + bSG_gC[(None, subtile_idx)], + ) + c_pipeline.producer_commit() + c_pipeline.producer_acquire() + self.epilog_sync_barrier.arrive_and_wait() + else: + # Convert to C type + acc_vec = tTR_rAcc.load() + acc_vec = acc_vec.to(self.c_dtype) + tTR_rC.store(acc_vec) + # Store C to global memory + cute.autovec_copy( + tTR_rC, tTR_gC[(None, None, None, subtile_idx)] + ) + # Async arrive accumulator buffer empty + with cute.arch.elect_one(): + acc_pipeline.consumer_release(acc_consumer_state) + acc_consumer_state.advance() + # Advance to next tile + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + + # Dealloc the tensor memory buffer + tmem.relinquish_alloc_permit() + self.epilog_sync_barrier.arrive_and_wait() + tmem.free(tmem_ptr) + if cutlass.const_expr(self.use_tma_store): + c_pipeline.producer_tail() + + # Idle warp + if warp_idx == self.idle_warp_id: + cute.arch.setmaxregister_decrease(self.num_regs_idle_warp) + + def epilog_gmem_copy_and_partition( + self, + tidx: cutlass.Int32, + atom: Union[cute.CopyAtom, cute.TiledCopy], + gC_mnl: cute.Tensor, + epi_tile: cute.Tile, + sC: cute.Tensor, + ) -> tuple[cute.CopyAtom, cute.Tensor, cute.Tensor]: + """ + Partitions source and destination tensors for a TMA store or SIMT store. + """ + if self.use_tma_store: + tma_atom_c, bSG_sC, bSG_gC, _, _ = ( + mixed_input_utils.epilog_gmem_copy_and_partition( + self.c_dtype, tidx, atom, None, gC_mnl, None, epi_tile, sC + ) + ) + return tma_atom_c, bSG_sC, bSG_gC + else: + _, _, _, simt_atom, tTR_gC = ( + mixed_input_utils.epilog_gmem_copy_and_partition( + self.c_dtype, tidx, None, atom, None, gC_mnl, epi_tile, sC + ) + ) + # (T2R, T2R_M, T2R_N) + tTR_rC = cute.make_rmem_tensor( + tTR_gC[(None, None, None, 0, 0, 0, 0, 0)].shape, self.c_dtype + ) + simt_atom = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), self.c_dtype) + return simt_atom, tTR_rC, tTR_gC + + @staticmethod + def _compute_stages_and_tmem_cols( + tiled_mma: cute.TiledMma, + mma_tiler_mnk: tuple[int, int, int], + cta_tile_shape_mnk: tuple[int, int, int], + epi_tile: cute.Tile, + a_dtype: type[cutlass.Numeric], + b_dtype: type[cutlass.Numeric], + c_dtype: type[cutlass.Numeric], + c_layout: utils.LayoutEnum, + transform_a_source: tcgen05.OperandSource, + scale_granularity_m: int, + scale_granularity_k: int, + smem_buffer_align_bytes: int, + use_tma_store: bool, + scale_mode: TransformMode, + ) -> tuple[int, int, int, int, int, int, int]: + """ + Compute pipeline stages and TMEM column allocation configurations. + + This method calculates the number of pipeline stages for different operations + (load2trans, trans2mma, accumulator, etc.) and determines TMEM column allocation + based on available memory resources and tile configuration. + + :param tiled_mma: The tiled MMA object defining the core computation. + :type tiled_mma: cute.TiledMma + :param mma_tiler_mnk: The shape (M, N, K) of the MMA tiler. + :type mma_tiler_mnk: tuple[int, int, int] + :param cta_tile_shape_mnk: The shape (M, N, K) of the CTA tile. + :type cta_tile_shape_mnk: tuple[int, int, int] + :param epi_tile: The epilogue tile shape. + :type epi_tile: cute.Tile + :param a_dtype: Data type of operand A. + :type a_dtype: type[cutlass.Numeric] + :param b_dtype: Data type of operand B. + :type b_dtype: type[cutlass.Numeric] + :param c_dtype: Data type of operand C. + :type c_dtype: type[cutlass.Numeric] + :param c_layout: Layout enum of operand C. + :type c_layout: utils.LayoutEnum + :param transform_a_source: The source of the transformed A tensor. + :type transform_a_source: tcgen05.OperandSource + :param scale_granularity_m: The granularity of the scale tensor along the M mode. + :type scale_granularity_m: int + :param scale_granularity_k: The granularity of the scale tensor along the K mode. + :type scale_granularity_k: int + :param smem_buffer_align_bytes: The alignment of the shared memory buffer. + :type smem_buffer_align_bytes: int + :param use_tma_store: Whether TMA store is enabled. + :type use_tma_store: bool + :param scale_mode: The transform mode. + :type scale_mode: TransformMode + + :return: A tuple containing the number of stages for: + (load2trans, scale_load2trans, transform2mma, accumulator, c, tmem_acc_cols, tmem_a_cols) + :rtype: tuple[int, int, int, int, int, int, int] + - num_load2trans_stage: Stages for load-to-transform A and B tensors pipeline + - num_scale_load2trans_stage: Stages for scale load-to-transform A tensor pipeline + - num_trans2mma_stage: Stages for transform-to-MMA pipeline + - num_acc_stage: Stages for accumulator-to-epilogue pipeline + - num_c_stage: Stages for epilogue-to-output C pipeline + - num_acc_tmem_cols: TMEM columns for accumulator + - num_a_tmem_cols: TMEM columns for transformed A tensor + """ + # Compute tmem columns required for accumulator + acc_shape = tiled_mma.partition_shape_C(mma_tiler_mnk[:2]) + tCtAcc_stage1 = tiled_mma.make_fragment_C(cute.append(acc_shape, 1)) + num_tmem_acc_col_per_stage = utils.get_num_tmem_alloc_cols(tCtAcc_stage1, True) + # Heuristic to decide the number of stages for accumulator + sm100_tmem_columns = cute.arch.get_max_tmem_alloc_cols("sm_100") + accumulator_stage_count = sm100_tmem_columns // num_tmem_acc_col_per_stage + if transform_a_source == tcgen05.OperandSource.TMEM: + if num_tmem_acc_col_per_stage < 128: + accumulator_stage_count = 3 + elif num_tmem_acc_col_per_stage < 256: + accumulator_stage_count = 2 + else: + accumulator_stage_count = 1 + # transformed A in 16bit, thus 1 tmem column could hold 2 elements + num_elts_per_tmem_col = 32 // tiled_mma.op.a_dtype.width + num_tmem_cols_a_per_stage = cute.round_up( + ( + cta_tile_shape_mnk[2] // num_elts_per_tmem_col + if transform_a_source == tcgen05.OperandSource.TMEM + else 0 + ), + 4, + ) + + c_stage_count = 2 if use_tma_store else 0 + c_smem_layout_staged_one = ( + sm100_utils.make_smem_layout_epi( + c_dtype, + c_layout, + epi_tile, + 1, + ) + if use_tma_store + else None + ) + c_bytes_per_stage = ( + cute.size_in_bytes(c_dtype, c_smem_layout_staged_one) + if use_tma_store + else 0 + ) + c_bytes = c_bytes_per_stage * c_stage_count + + smem_capacity = utils.get_smem_capacity_in_bytes("sm_100") + bytes_per_pipeline_stage = 16 + if scale_mode == TransformMode.ConvertOnly: + scale_load2trans_stage_count = 0 + a_scale_bytes_per_stage = 0 + else: + # Ensure we have 2 buffers for scale tiles needed for 1 CTA tile + a_scale_k_mode = max(cta_tile_shape_mnk[2] // scale_granularity_k, 1) + a_scale_m_mode = max(cta_tile_shape_mnk[0] // scale_granularity_m, 1) + scale_load2trans_stage_count = 4 + a_scale_bytes_per_stage = cute.round_up( + cute.size_in_bytes( + tiled_mma.op.a_dtype, + cute.make_layout((a_scale_m_mode, a_scale_k_mode)), + ), + smem_buffer_align_bytes, + ) + a_scale_bytes = ( + a_scale_bytes_per_stage + bytes_per_pipeline_stage + ) * scale_load2trans_stage_count + carveout_smem_bytes = ( + bytes_per_pipeline_stage * accumulator_stage_count + a_scale_bytes + c_bytes + ) + + # Compute transform stages if A is in TMEM + num_tmem_acc_cols = cute.round_up( + accumulator_stage_count * num_tmem_acc_col_per_stage, 4 + ) + + transform2mma_stage_count_a_source_tmem_potential = ( + (sm100_tmem_columns - num_tmem_acc_cols) // num_tmem_cols_a_per_stage + if transform_a_source == tcgen05.OperandSource.TMEM + else -1 + ) + if ( + transform_a_source == tcgen05.OperandSource.TMEM + and transform2mma_stage_count_a_source_tmem_potential <= 0 + ): + raise ValueError("Not enough TMEM capacity for selected tile size") + a_load_bytes_per_stage = cute.round_up( + cute.size_in_bytes( + a_dtype, + cute.make_layout((cta_tile_shape_mnk[0], cta_tile_shape_mnk[2])), + ), + smem_buffer_align_bytes, + ) + b_load_bytes_per_stage = cute.round_up( + cute.size_in_bytes( + b_dtype, + cute.make_layout( + ( + cta_tile_shape_mnk[1] // cute.size(tiled_mma.thr_id), + cta_tile_shape_mnk[2], + ) + ), + ), + smem_buffer_align_bytes, + ) + ab_load_bytes_per_stage = ( + a_load_bytes_per_stage + + b_load_bytes_per_stage + + 2 * bytes_per_pipeline_stage + ) + a_transform_bytes_per_stage = ( + cute.round_up( + cute.size_in_bytes( + tiled_mma.op.a_dtype, + cute.make_layout((cta_tile_shape_mnk[0], cta_tile_shape_mnk[2])), + ), + smem_buffer_align_bytes, + ) + if transform_a_source == tcgen05.OperandSource.SMEM + else 0 + ) + + a_transform_bytes_per_stage = ( + a_transform_bytes_per_stage + bytes_per_pipeline_stage + ) + transform2mma_stage_count_a_source_smem_potential = ( + smem_capacity - carveout_smem_bytes + ) // (ab_load_bytes_per_stage + a_transform_bytes_per_stage) + transform2mma_stage_count = ( + min( + transform2mma_stage_count_a_source_tmem_potential, + transform2mma_stage_count_a_source_smem_potential, + ) + if transform_a_source == tcgen05.OperandSource.TMEM + else transform2mma_stage_count_a_source_smem_potential + ) + load2transform_stage_count = ( + smem_capacity + - carveout_smem_bytes + - (transform2mma_stage_count * a_transform_bytes_per_stage) + ) // ab_load_bytes_per_stage + if ( + load2transform_stage_count < 2 + or transform2mma_stage_count < 2 + or accumulator_stage_count < 1 + ): + raise ValueError("Not enough SMEM or TMEM capacity for selected tile size") + num_tmem_a_cols = transform2mma_stage_count * num_tmem_cols_a_per_stage + # Check if we can increase c_stage_count with leftover smem + if use_tma_store: + c_stage_count += ( + smem_capacity + - load2transform_stage_count * ab_load_bytes_per_stage + - transform2mma_stage_count * a_transform_bytes_per_stage + - scale_load2trans_stage_count * a_scale_bytes_per_stage + - c_bytes + ) // c_bytes_per_stage + + return ( + load2transform_stage_count, + scale_load2trans_stage_count, + transform2mma_stage_count, + accumulator_stage_count, + c_stage_count, + num_tmem_acc_cols, + num_tmem_a_cols, + ) + + @staticmethod + def _compute_grid( + c: cute.Tensor, + cta_tile_shape_mnk: tuple[int, int, int], + cluster_shape_mn: tuple[int, int], + max_active_clusters: cutlass.Constexpr, + ) -> tuple[utils.PersistentTileSchedulerParams, tuple[int, int, int]]: + """ + Use persistent tile scheduler to compute the grid size for the output tensor C. + """ + c_shape = cute.slice_(cta_tile_shape_mnk, (None, None, 0)) + gc = cute.zipped_divide(c, tiler=c_shape) + num_ctas_mnl = gc[(0, (None, None, None))].shape + cluster_shape_mnl = (*cluster_shape_mn, 1) + + tile_sched_params = utils.PersistentTileSchedulerParams( + num_ctas_mnl, cluster_shape_mnl + ) + grid = utils.StaticPersistentTileScheduler.get_grid_shape( + tile_sched_params, max_active_clusters + ) + + return tile_sched_params, grid + + def is_valid_epilog_store_option( + m: int, + n: int, + mma_tiler_mn: tuple[int, int], + use_tma_store: bool, + use_2cta_instrs: bool, + ) -> bool: + """ + Check if the epilogue store option is valid for the given problem size. + """ + cta_tile_shape_mn = ( + mma_tiler_mn[0] // (2 if use_2cta_instrs else 1), + mma_tiler_mn[1], + ) + # No OOB tile support when TMA store is disabled + if not use_tma_store: + if not (m % cta_tile_shape_mn[0] == 0 and n % cta_tile_shape_mn[1] == 0): + return False + return True + + def can_implement( + mnkl: tuple[int, int, int, int], + a_dtype: type[cutlass.Numeric], + b_dtype: type[cutlass.Numeric], + c_dtype: type[cutlass.Numeric], + a_major: str, + b_major: str, + c_major: str, + scale_granularity_m: int, + scale_granularity_k: int, + mma_tiler: tuple[int, int, int], + cluster_shape_mn: tuple[int, int], + use_2cta_instrs: bool, + use_tma_store: bool, + ) -> bool: + """ + Check if the kernel can be implemented for the given tensor shapes and data types. + """ + m, n, k, l = mnkl + + if not mixed_input_utils.is_valid_mma_tiler_and_cluster_shape( + mma_tiler, cluster_shape_mn, use_2cta_instrs + ): + return False + if not mixed_input_utils.is_valid_scale_granularity( + scale_granularity_m, scale_granularity_k, a_dtype, k, mma_tiler[2] + ): + return False + if not mixed_input_utils.is_valid_tensor_alignment( + m, + n, + k, + a_dtype, + b_dtype, + c_dtype, + b_dtype, + a_major, + b_major, + c_major, + mma_tiler, + use_2cta_instrs, + cluster_shape_mn, + scale_granularity_m, + scale_granularity_k, + ): + return False + if not MixedInputGemmKernel.is_valid_epilog_store_option( + m, n, mma_tiler[:2], use_tma_store, use_2cta_instrs + ): + return False + return True + + +def run( + mnkl: tuple[int, int, int, int], + scale_granularity_m: int, + scale_granularity_k: int, + a_dtype: type[cutlass.Numeric], + b_dtype: type[cutlass.Numeric], + c_dtype: type[cutlass.Numeric], + acc_dtype: type[cutlass.Numeric], + a_major: str, + b_major: str, + c_major: str, + mma_tiler_mnk: tuple[int, int, int], + cluster_shape_mn: tuple[int, int], + use_2cta_instrs: bool, + use_tma_store: bool, + tolerance: float, + warmup_iterations: int = 0, + iterations: int = 1, + skip_ref_check: bool = False, + use_cold_l2: bool = False, + **kwargs, +) -> None: + """ + Run the mixed-input GEMM kernel with specified parameters. + + This function creates tensors, validates parameters, executes the kernel, + optionally compares results with a reference implementation and reports + kernel execution time. + """ + m, n, k, l = mnkl + import torch + import cutlass.torch as cutlass_torch + + if not torch.cuda.is_available(): + raise ValueError("CUDA is not available") + + # Check if given configuration is supported + if not MixedInputGemmKernel.can_implement( + mnkl, + a_dtype, + b_dtype, + c_dtype, + a_major, + b_major, + c_major, + scale_granularity_m, + scale_granularity_k, + mma_tiler_mnk, + cluster_shape_mn, + use_2cta_instrs, + use_tma_store, + ): + raise ValueError("GEMM configuration not supported") + + # Get current CUDA stream from PyTorch + torch_stream = torch.cuda.current_stream() + # Get the raw stream pointer as a CUstream + current_stream = cuda.CUstream(torch_stream.cuda_stream) + shuffle_a = mixed_input_utils.is_shuffle_a( + a_major, k, a_dtype, b_dtype, scale_granularity_k + ) + # shuffle is supported since CUDA 13.1 + shuffle_supported = False + try: + from cutlass import CUDA_VERSION + + if CUDA_VERSION.major > 13 or ( + CUDA_VERSION.major == 13 and CUDA_VERSION.minor >= 1 + ): + shuffle_supported = True + except ImportError: + pass + + shuffle_a = shuffle_a and shuffle_supported + mixed_input_gemm = MixedInputGemmKernel( + scale_granularity_m, + scale_granularity_k, + acc_dtype, + use_2cta_instrs, + mma_tiler_mnk, + cluster_shape_mn, + use_tma_store, + shuffle_a, + ) + ( + a_tensor, + a_scale_tensor, + b_tensor, + c_tensor, + a_torch_cpu, + a_scale_torch_cpu, + b_torch_cpu, + c_torch_gpu, + ) = create_tensors( + l, + m, + n, + k, + a_major, + b_major, + c_major, + a_dtype, + b_dtype, + c_dtype, + shuffle_a, + scale_granularity_m, + scale_granularity_k, + ) + + max_active_clusters = utils.HardwareInfo().get_max_active_clusters( + cluster_shape_mn[0] * cluster_shape_mn[1], + ) + # try to check CUDA version to decide the opt level + try: + from cutlass import CUDA_VERSION + + opt_level = 3 if (CUDA_VERSION.major == 13 and CUDA_VERSION.minor < 1) else 2 + except ImportError: + opt_level = 3 + compiled_kernel = cute.compile( + mixed_input_gemm, + a_tensor, + a_scale_tensor, + b_tensor, + c_tensor, + max_active_clusters, + current_stream, + options=f"--opt-level {opt_level}", + ) + + if not skip_ref_check: + compiled_kernel( + a_tensor, + a_scale_tensor, + b_tensor, + c_tensor, + current_stream, + ) + run_ref_and_compare( + a_torch_cpu, b_torch_cpu, a_scale_torch_cpu, c_torch_gpu, c_dtype, tolerance + ) + + # Early return if no performance measurement is needed + if iterations <= 0: + return + + def generate_tensors(): + ( + a_tensor, + a_scale_tensor, + b_tensor, + c_tensor, + a_torch_cpu, + a_scale_torch_cpu, + b_torch_cpu, + c_torch_gpu, + ) = create_tensors( + l, + m, + n, + k, + a_major, + b_major, + c_major, + a_dtype, + b_dtype, + c_dtype, + shuffle_a, + scale_granularity_m, + scale_granularity_k, + ) + return testing.JitArguments( + a_tensor, a_scale_tensor, b_tensor, c_tensor, current_stream + ) + + workspace_count = 1 + if use_cold_l2: + one_workspace_bytes = ( + a_torch_cpu.numel() * a_torch_cpu.element_size() + + b_torch_cpu.numel() * b_torch_cpu.element_size() + + c_torch_gpu.numel() * c_torch_gpu.element_size() + + a_scale_torch_cpu.numel() * a_scale_torch_cpu.element_size() + if a_scale_torch_cpu is not None + else 0 + ) + workspace_count = testing.get_workspace_count( + one_workspace_bytes, warmup_iterations, iterations + ) + + exec_time = testing.benchmark( + compiled_kernel, + workspace_generator=generate_tensors, + workspace_count=workspace_count, + stream=current_stream, + warmup_iterations=warmup_iterations, + iterations=iterations, + ) + + return exec_time # Return execution time in microseconds + + +if __name__ == "__main__": + + def parse_comma_separated_ints(s: str) -> tuple[int, ...]: + try: + return tuple(int(x.strip()) for x in s.split(",")) + except ValueError: + raise argparse.ArgumentTypeError( + "Invalid format. Expected comma-separated integers." + ) + + parser = argparse.ArgumentParser() + parser.add_argument( + "--mnkl", type=parse_comma_separated_ints, default=(128, 128, 128, 1) + ) + parser.add_argument( + "--mma_tiler_mnk", type=parse_comma_separated_ints, default=(128, 128, 128) + ) + parser.add_argument( + "--cluster_shape_mn", type=parse_comma_separated_ints, default=(1, 1) + ) + parser.add_argument( + "--use_2cta_instrs", + action="store_true", + help="Enable 2CTA MMA instructions feature", + ) + parser.add_argument( + "--a_dtype", + type=cutlass.dtype, + default=cutlass.Int4, + choices=[cutlass.Int8, cutlass.Uint8, cutlass.Int4], + ) + parser.add_argument( + "--b_dtype", + type=cutlass.dtype, + default=cutlass.BFloat16, + choices=[cutlass.BFloat16, cutlass.Float16], + ) + parser.add_argument("--c_dtype", type=cutlass.dtype, default=cutlass.BFloat16) + parser.add_argument("--acc_dtype", type=cutlass.dtype, default=cutlass.Float32) + parser.add_argument("--a_major", choices=["k", "m"], type=str, default="m") + parser.add_argument("--b_major", choices=["k", "n"], type=str, default="k") + parser.add_argument("--c_major", choices=["n", "m"], type=str, default="n") + parser.add_argument( + "--scale_granularity_m", + type=int, + default=1, + help="Scale granularity along M dimension.", + ) + parser.add_argument( + "--scale_granularity_k", + type=int, + default=128, + help="Scale granularity along K dimension.", + ) + parser.add_argument( + "--use_tma_store", action="store_true", help="Use tma store or not" + ) + parser.add_argument( + "--tolerance", type=float, default=1e-01, help="Tolerance for validation" + ) + parser.add_argument( + "--warmup_iterations", type=int, default=0, help="Warmup iterations" + ) + parser.add_argument( + "--iterations", + type=int, + default=1, + help="Number of iterations to run the kernel", + ) + parser.add_argument( + "--skip_ref_check", action="store_true", help="Skip reference checking" + ) + args = parser.parse_args() + + run( + args.mnkl, + args.scale_granularity_m, + args.scale_granularity_k, + args.a_dtype, + args.b_dtype, + args.c_dtype, + args.acc_dtype, + args.a_major, + args.b_major, + args.c_major, + args.mma_tiler_mnk, + args.cluster_shape_mn, + args.use_2cta_instrs, + args.use_tma_store, + args.tolerance, + args.warmup_iterations, + args.iterations, + args.skip_ref_check, + ) + print("PASS") diff --git a/examples/python/CuTeDSL/blackwell/mixed_input_gemm/mixed_input_host_utils.py b/examples/python/CuTeDSL/blackwell/mixed_input_gemm/mixed_input_host_utils.py new file mode 100644 index 00000000..107dc5a7 --- /dev/null +++ b/examples/python/CuTeDSL/blackwell/mixed_input_gemm/mixed_input_host_utils.py @@ -0,0 +1,506 @@ +# 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. + +from typing import Optional + +import torch + +import cutlass +import cutlass.cute as cute +import cutlass.torch as cutlass_torch +import cutlass.utils.mixed_input_helpers as mixed_input_utils +from cutlass.cute.runtime import from_dlpack + +""" +This file contains common host-side utilities for mixed-input GEMM. +""" + + +def create_cumsum_tensor( + num_groups: int, + fused_n: int, + alignment: int, + uniform_distribution: bool = False, +) -> tuple[cute.Tensor, torch.Tensor]: + """ + Create a tensor of shape (num_groups + 1) recording the cumulative sum of the elements in each group. + """ + assert fused_n % alignment == 0, "fused_n must be divisible by alignment" + if uniform_distribution: + # keep a uniform distribution for debug and performance collection + group_counts = torch.tensor([fused_n // num_groups] * num_groups) + else: + # sample group sizes with equal probability for each group + probs = torch.ones(num_groups) / num_groups + group_sizes = torch.multinomial(probs, fused_n // alignment, replacement=True) + group_counts = torch.bincount(group_sizes, minlength=num_groups) * alignment + print(group_counts.tolist()) + + # Create cumulative sum + cumsum_torch = torch.cat([torch.tensor([0]), group_counts.cumsum(0)]) + print(cumsum_torch.tolist()) + + cumsum_tensor, _ = cutlass_torch.cute_tensor_like( + cumsum_torch, cutlass.Int32, is_dynamic_layout=False + ) + + return cumsum_tensor, cumsum_torch.to("cpu") + + +def create_i4_tensor_and_scale( + l: int, + m: int, + k: int, + is_m_major: bool, + dtype: type[cutlass.Numeric], + shuffle_a: bool, + scale_granularity_m: int, + scale_granularity_k: int, + is_dynamic_layout: bool = True, + init_config: tuple = ( + cutlass_torch.TensorInitType.RANDOM, + cutlass_torch.RandomInitConfig(min_val=-7, max_val=6), + ), + divisibility: int = 16, + transformed_dtype: Optional[type[cutlass.Numeric]] = None, +) -> tuple[ + cute.Tensor, + torch.Tensor, + torch.Tensor, + cute.Tensor, + torch.Tensor, + torch.Tensor, +]: + """ + Create quantized 4-bit tensor and corresponding scale tensor. + """ + lb_4b = -8 if dtype == cutlass.Int4 else 0 + up_4b = 7 if dtype == cutlass.Int4 else 15 + if not ( + init_config[0] == cutlass_torch.TensorInitType.RANDOM + or init_config[0] == cutlass_torch.TensorInitType.SCALAR + ): + raise ValueError( + "Only random and scalar initialization is supported for 4bit data type" + ) + + # Construct reference tensor in f32 + ref_fp32 = cutlass_torch.matrix(l, m, k, is_m_major, cutlass.Float32, *init_config) + # Generate scale data and perform quantization + num_scales = k // scale_granularity_k + ref = ref_fp32.to(dtype=cutlass_torch.dtype(transformed_dtype)).reshape( + m, num_scales, scale_granularity_k, l + ) + # Get elements with maximum absolute value to compute scaling factors + a_max = ( + torch.maximum(ref / up_4b, ref / lb_4b) + if dtype == cutlass.Int4 + else ref / up_4b + ) + a_scales, _ = torch.max(a_max, dim=2, keepdim=True) + a_scale_inv = torch.where(a_scales == 0, 0, 1 / a_scales) + a_quant = ref * a_scale_inv + # Convert values to integer to avoid computation errors + a_quant = a_quant.to(dtype=torch.int32).reshape((m, k, l)).to(dtype=torch.float32) + # Construct cute scale tensor + a_scales = a_scales.random_(-3, 3).reshape((m, num_scales, l)) + # Scale tensor is always m-major + a_scales = a_scales.permute(2, 1, 0).contiguous().permute(2, 1, 0).to(device="cuda") + if shuffle_a: + # shuffle within each group of 8 elements + perm = torch.tensor([0, 2, 1, 3, 4, 6, 5, 7], device=a_quant.device) + a_shuffled = ( + a_quant.view(m, k // 8, 8, l)[:, :, perm, :] + .reshape(a_quant.shape) + .permute(2, 0, 1) + .contiguous() + .permute(1, 2, 0) + ) + # Construct A quantized tensor + cute_a_quant_tensor, torch_a_quant_tensor = cutlass_torch.cute_tensor_like( + a_shuffled, + dtype, + is_dynamic_layout=is_dynamic_layout, + assumed_align=divisibility, + ) + else: + # Construct A quantized tensor + cute_a_quant_tensor, torch_a_quant_tensor = cutlass_torch.cute_tensor_like( + a_quant, + dtype, + is_dynamic_layout=is_dynamic_layout, + assumed_align=divisibility, + ) + cute_scale_tensor = from_dlpack(a_scales, assumed_align=divisibility) + for i, stride in enumerate(a_scales.stride()): + if stride == 1: + leading_dim = i + break + if is_dynamic_layout: + cute_scale_tensor = cute_scale_tensor.mark_layout_dynamic( + leading_dim=leading_dim + ) + + return ( + cute_a_quant_tensor, + torch_a_quant_tensor, + a_quant.to("cpu"), + cute_scale_tensor, + a_scales, + a_scales.to("cpu"), + ) + + +def create_tensor_a( + l: int, + m: int, + k: int, + a_major: str, + a_dtype: type[cutlass.Numeric], + shuffle_a: bool, + scale_granularity_m: int = 0, + scale_granularity_k: int = 0, + transformed_dtype: Optional[type[cutlass.Numeric]] = None, +) -> tuple[cute.Tensor, Optional[cute.Tensor], torch.Tensor, Optional[torch.Tensor]]: + """ + Create tensor A and scale tensor. + """ + a_scale_tensor = None + a_scale_torch_cpu = None + if a_dtype in (cutlass.Int4,): + ( + a_tensor, + a_torch_gpu, + a_torch_cpu, + a_scale_tensor, + a_scale_torch_gpu, + a_scale_torch_cpu, + ) = create_i4_tensor_and_scale( + l, + m, + k, + a_major == "m", + a_dtype, + shuffle_a, + scale_granularity_m, + scale_granularity_k, + divisibility=mixed_input_utils.get_divisibility(m if a_major == "m" else k), + transformed_dtype=transformed_dtype, + ) + else: + a_torch_cpu = cutlass_torch.matrix( + l, + m, + k, + a_major == "m", + a_dtype, + ) + a_tensor, _ = cutlass_torch.cute_tensor_like( + a_torch_cpu, + a_dtype, + is_dynamic_layout=True, + assumed_align=mixed_input_utils.get_divisibility( + m if a_major == "m" else k + ), + ) + return a_tensor, a_scale_tensor, a_torch_cpu, a_scale_torch_cpu + + +def create_tensors_for_contiguous_grouped_mixed_input_gemm( + l: int, + m: int, + n: int, + k: int, + a_major: str, + b_major: str, + c_major: str, + a_dtype: type[cutlass.Numeric], + b_dtype: type[cutlass.Numeric], + c_dtype: type[cutlass.Numeric], + shuffle_a: bool = False, + scale_granularity_m: int = 0, + scale_granularity_k: int = 0, + uniform_group_sizes: bool = False, +) -> tuple: + """ + Create all input and output tensors for the contiguous grouped mixed-input GEMM. + """ + a_tensor, a_scale_tensor, a_torch_cpu, a_scale_torch_cpu = create_tensor_a( + l, + m, + k, + a_major, + a_dtype, + shuffle_a, + scale_granularity_m, + scale_granularity_k, + b_dtype, + ) + + # In GROUP mode, l specifies the number of groups. We'll fuse group into the n mode for tensor B and C. + # Batch mode will be set to 1. + num_groups = l + fused_n = n * num_groups + b_torch_cpu = cutlass_torch.matrix( + 1, # batch=1 + fused_n, + k, + b_major == "n", + b_dtype, + cutlass_torch.TensorInitType.RANDOM, + cutlass_torch.RandomInitConfig(min_val=-10, max_val=10), + ) + b_tensor, _ = cutlass_torch.cute_tensor_like( + b_torch_cpu, + b_dtype, + is_dynamic_layout=True, + assumed_align=mixed_input_utils.get_divisibility(n if b_major == "n" else k), + ) + + c_torch_cpu = cutlass_torch.matrix( + 1, # batch=1 + m, + fused_n, + c_major == "m", + c_dtype, + ) + c_tensor, c_torch_gpu = cutlass_torch.cute_tensor_like( + c_torch_cpu, + c_dtype, + is_dynamic_layout=True, + assumed_align=mixed_input_utils.get_divisibility(m if c_major == "m" else n), + ) + c_tensor = c_tensor.mark_compact_shape_dynamic( + mode=(0 if c_major == "m" else 1), + stride_order=(2, 1, 0) if c_major == "m" else (2, 0, 1), + divisibility=mixed_input_utils.get_divisibility(m if c_major == "m" else n), + ) + # We need to ensure mode N satisfies 16B alignment for each group + alignment_n = 16 * 8 // b_dtype.width + cumsum_tensor, cumsum_torch = create_cumsum_tensor( + num_groups, fused_n, alignment_n, uniform_distribution=uniform_group_sizes + ) + + return ( + a_tensor, + a_scale_tensor, + b_tensor, + cumsum_tensor, + c_tensor, + a_torch_cpu, + a_scale_torch_cpu, + b_torch_cpu, + cumsum_torch, + c_torch_gpu, + ) + + +def create_tensors_for_batched_mixed_input_gemm( + l: int, + m: int, + n: int, + k: int, + a_major: str, + b_major: str, + c_major: str, + a_dtype: type[cutlass.Numeric], + b_dtype: type[cutlass.Numeric], + c_dtype: type[cutlass.Numeric], + shuffle_a: bool = False, + scale_granularity_m: int = 0, + scale_granularity_k: int = 0, +) -> tuple: + """ + Create all input and output tensors for the batched mixed-input GEMM. + """ + torch.manual_seed(2025) + + a_tensor, a_scale_tensor, a_torch_cpu, a_scale_torch_cpu = create_tensor_a( + l, + m, + k, + a_major, + a_dtype, + shuffle_a, + scale_granularity_m, + scale_granularity_k, + b_dtype, + ) + + b_torch_cpu = cutlass_torch.matrix( + l, + n, + k, + b_major == "n", + b_dtype, + cutlass_torch.TensorInitType.RANDOM, + cutlass_torch.RandomInitConfig(min_val=-10, max_val=10), + ) + c_torch_cpu = cutlass_torch.matrix( + l, + m, + n, + c_major == "m", + c_dtype, + ) + + b_tensor, _ = cutlass_torch.cute_tensor_like( + b_torch_cpu, + b_dtype, + is_dynamic_layout=True, + assumed_align=mixed_input_utils.get_divisibility(n if b_major == "n" else k), + ) + c_tensor, c_torch_gpu = cutlass_torch.cute_tensor_like( + c_torch_cpu, + c_dtype, + is_dynamic_layout=True, + assumed_align=mixed_input_utils.get_divisibility(m if c_major == "m" else n), + ) + c_tensor = c_tensor.mark_compact_shape_dynamic( + mode=(0 if c_major == "m" else 1), + stride_order=(2, 1, 0) if c_major == "m" else (2, 0, 1), + divisibility=mixed_input_utils.get_divisibility(m if c_major == "m" else n), + ) + + return ( + a_tensor, + a_scale_tensor, + b_tensor, + c_tensor, + a_torch_cpu, + a_scale_torch_cpu, + b_torch_cpu, + c_torch_gpu, + ) + + +def run_contiguous_grouped_ref_and_compare( + a_torch_cpu: torch.Tensor, + b_torch_cpu: torch.Tensor, + a_scale_torch_cpu: Optional[torch.Tensor], + cumsum_torch_cpu: torch.Tensor, + c_torch_gpu: torch.Tensor, + c_dtype: type[cutlass.Numeric], + tolerance: float, +) -> None: + """ + Compare kernel result with reference computation. + """ + kernel_result = c_torch_gpu.cpu() + assert kernel_result.shape[2] == 1, "batch mode must be 1" + kernel_result = kernel_result.reshape( + kernel_result.shape[0], kernel_result.shape[1] + ) + # Compute reference result + a_for_gemm = a_torch_cpu + if a_scale_torch_cpu is not None: + scale_shape = a_scale_torch_cpu.shape + a_shape = a_torch_cpu.shape + a_scale_torch_cpu = a_scale_torch_cpu.to(dtype=torch.float32).reshape( + scale_shape[0], scale_shape[1], 1, scale_shape[2] + ) + a_torch_cpu = a_torch_cpu.to(dtype=torch.float32).reshape( + a_torch_cpu.shape[0], scale_shape[1], -1, a_torch_cpu.shape[2] + ) + a_for_gemm = (a_torch_cpu * a_scale_torch_cpu).reshape(a_shape) + # A in (m, k, l), b in (n, k), c in (m, n) + assert cumsum_torch_cpu.shape[0] == a_for_gemm.shape[-1] + 1, ( + "cumsum tensor must have one more element than a_for_gemm" + ) + assert b_torch_cpu.shape[2] == 1, ( + "b_torch_cpu must have a singleton dimension in the last position" + ) + prev_idx = 0 + ref = torch.zeros((a_for_gemm.shape[0], b_torch_cpu.shape[0]), dtype=torch.float32) + for group_idx in range(1, cumsum_torch_cpu.shape[0]): + # No computation for current group + if cumsum_torch_cpu[group_idx] == prev_idx: + continue + # Get A slice for current group + sliced_a = a_for_gemm[:, :, group_idx - 1] + # Get B slice for current group + sliced_b = b_torch_cpu[prev_idx : cumsum_torch_cpu[group_idx], :, 0] + sliced_ref = torch.einsum( + "mk,nk->mn", + sliced_a.to(dtype=torch.float32), + sliced_b.to(dtype=torch.float32), + ) + ref[:, prev_idx : cumsum_torch_cpu[group_idx]] = sliced_ref + prev_idx = cumsum_torch_cpu[group_idx] + # Convert ref to c_dtype + _, ref_torch_gpu = cutlass_torch.cute_tensor_like( + ref, c_dtype, is_dynamic_layout=True, assumed_align=16 + ) + ref_result = ref_torch_gpu.cpu() + + # Assert close results + torch.testing.assert_close(kernel_result, ref_result, atol=tolerance, rtol=1e-05) + + +def run_batched_mixed_input_ref_and_compare( + a_torch_cpu: torch.Tensor, + b_torch_cpu: torch.Tensor, + a_scale_torch_cpu: Optional[torch.Tensor], + c_torch_gpu: torch.Tensor, + c_dtype: type[cutlass.Numeric], + tolerance: float, +) -> None: + """ + Compare kernel result with reference computation. + """ + kernel_result = c_torch_gpu.cpu() + # Compute reference result + if a_scale_torch_cpu is not None: + scale_shape = a_scale_torch_cpu.shape + a_shape = a_torch_cpu.shape + a_scale_torch_cpu = a_scale_torch_cpu.to(dtype=torch.float32).reshape( + scale_shape[0], scale_shape[1], 1, scale_shape[2] + ) + a_torch_cpu = a_torch_cpu.to(dtype=torch.float32).reshape( + a_torch_cpu.shape[0], scale_shape[1], -1, a_torch_cpu.shape[2] + ) + a_dequant = a_torch_cpu * a_scale_torch_cpu + ref = torch.einsum( + "mkl,nkl->mnl", + a_dequant.reshape(a_shape), + b_torch_cpu.to(dtype=torch.float32), + ) + else: + ref = torch.einsum( + "mkl,nkl->mnl", + a_torch_cpu.to(dtype=torch.float32), + b_torch_cpu.to(dtype=torch.float32), + ) + # Convert ref to c_dtype + _, ref_torch_gpu = cutlass_torch.cute_tensor_like( + ref, c_dtype, is_dynamic_layout=True, assumed_align=16 + ) + ref_result = ref_torch_gpu.cpu() + + # Assert close results + torch.testing.assert_close(kernel_result, ref_result, atol=tolerance, rtol=1e-05) diff --git a/examples/python/CuTeDSL/blackwell/mla/mla_decode_fp16.py b/examples/python/CuTeDSL/blackwell/mla/mla_decode_fp16.py new file mode 100644 index 00000000..7eba992b --- /dev/null +++ b/examples/python/CuTeDSL/blackwell/mla/mla_decode_fp16.py @@ -0,0 +1,4373 @@ +# 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. + +import os +import sys +import argparse +import math +from typing import Type, Tuple, Optional +from types import SimpleNamespace + +import cuda.bindings.driver as cuda + +import cutlass +import cutlass.cute as cute +import cutlass.cute.testing as testing +import cutlass.cute.nvgpu.tcgen05 as tcgen05 +from cutlass.cute.nvgpu.tcgen05 import OperandMajorMode +import cutlass.cute.nvgpu.cpasync as cpasync +import cutlass.utils as utils +import cutlass.pipeline as pipeline +from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait +import cutlass.utils.blackwell_helpers as sm100_utils +from cutlass.cute.runtime import from_dlpack +from cutlass.base_dsl.arch import Arch +from cutlass.cutlass_dsl import BaseDSL + +if __name__ == "__main__": + current_dir = os.path.dirname(os.path.abspath(__file__)) + sys.path.insert(0, os.path.join(current_dir, "../..")) + +from blackwell.mla.mla_helpers import ( + ceil_div, + MAX_SPLITS, + LOG2_E, + MLAStaticTileScheduler, + MLAStaticTileSchedulerParams, + create_mla_static_tile_scheduler, + create_mla_static_tile_scheduler_params, +) + +""" +A Multi-Head Latent Attention (MLA) example with FP16 data type for the NVIDIA Blackwell SM100 architecture using CUTE DSL + +This example demonstrates an implementation of inference of multi-head latent attention using a TMA + Blackwell +SM100 TensorCore warp-specialized persistent kernel. The implementation integrates the (Qc + Qr)*(Kc + Kr)^T +matrix multiplication, softmax normalization, and softmax((Qc + Qr)*(Kc + Kr)^T)*Vc into a single kernel. +The kernel provides support for page table storage and variable-length KV cache sequences. It implements KV splitting +functionality to minimize latency when processing long KV sequences. + +The kernel implements key optimizations including: +- Warp specialization for different computation phases (load, MMA, softmax, correction, epilogue) +- Pipeline stages between different warps for overlapping computation and memory access +- Support for different precision data types +- Two sub-kernels (split KV kernel and reduction kernel) that enable split KV processing + +To run this example: + +.. code-block:: bash + + python examples/blackwell/mla_fp16.py \ + --batch_size 4 --latent_dim 512 --rope_dim 64 \ + --num_heads 128 --seq_len_q 1 --seq_len_k 1024 \ + --in_dtype Float16 --out_dtype Float16 \ + --acc_dtype Float32 --lse_dtype Float32 \ + --is_var_seq --is_var_split_kv \ + --is_persistent + +The above example runs Multi-Head Latent Attention (MLA) with the following configuration: +- Batch size: 4 +- Sequence length of Q: 1 +- Sequence length of K: 1024 +- Latent dimension: 512 +- RoPE dimension: 64 +- Number of heads: 128 +- Data types: Float16 (input), Float16 (output), Float32 (accumulation and LSE) + +It utilizes page table storage for the KV cache and enables both variable-length KV cache sequences +and variable split KV processing with persistent scheduling. + +To collect performance with NCU profiler: + +.. code-block:: bash + + ncu python examples/blackwell/mla_fp16.py \ + --batch_size 4 --latent_dim 512 --rope_dim 64 \ + --num_heads 128 --seq_len_q 1 --seq_len_k 1024 \ + --in_dtype Float16 --out_dtype Float16 \ + --acc_dtype Float32 --lse_dtype Float32 \ + --is_var_seq --is_var_split_kv \ + --is_persistent --warmup_iterations 3 \ + --iterations 10 --skip_ref_check + +Constraints for this example: +* Data type requirements: + - Input/output: Float16 + - Accumulation and LSE: Float32 +* Fixed architecture parameters: + - Number of attention heads: 128 + - Latent dimension: 512 + - RoPE dimension: 64 +* Input query modes should be (NumHeads, LatentDim/RopeDim, SeqLenQ, BatchSize) +* Input kv latent/rope modes should be (SeqLenK, LatentDim/RopeDim, BatchSize) +* Query sequence length must be 1-4 +* Only supports 2-CTA instructions +* Variable sequence length requires page table storage enabled +""" + + +class BlackwellMultiHeadLatentAttentionForwardFP16: + def __init__( + self, + acc_dtype: Type[cutlass.Numeric], + lse_dtype: Type[cutlass.Numeric], + mma_qk_tiler_mn: Tuple[int, int], + mma_pv_tiler_mn: Tuple[int, int], + max_active_clusters: int, + page_size: int, + skip_correction_threshold: float, + is_persistent: bool, + is_var_seq: bool, + is_var_split_kv: bool, + ): + """Initializes the configuration for a Blackwell Multi-Head Latent Attention (MLA) kernel. + + :param acc_dtype: Data type for accumulation S and O + :type acc_dtype: Type[cutlass.Numeric] + :param lse_dtype: Data type for output LSE + :type lse_dtype: Type[cutlass.Numeric] + :param mma_s_tiler: The (H, K) tile shape of the MMA instruction for S + :type mma_s_tiler: Tuple[int, int] + :param mma_p_tiler: The (H, D) tile shape of the MMA instruction for P + :type mma_p_tiler: Tuple[int, int] + :param max_active_clusters: Maximum number of active clusters + :type max_active_clusters: int + :param page_size: The page size of the page table + :type page_size: int + :param skip_correction_threshold: Threshold to skip correction + :type skip_correction_threshold: float + :param is_persistent: Whether to use persistent kernel mode + :type is_persistent: bool + :param is_var_seq: Whether to use variable sequence length + :type is_var_seq: bool + :param is_var_split_kv: Whether to use variable split KV + :type is_var_split_kv: bool + """ + + self.latent_dim = 512 + self.rope_dim = 64 + self.acc_dtype = acc_dtype + self.lse_dtype = lse_dtype + self.mma_qk_tiler_mn = mma_qk_tiler_mn + self.mma_pv_tiler_mn = mma_pv_tiler_mn + self.max_active_clusters = max_active_clusters + self.skip_correction_threshold = skip_correction_threshold + self.is_persistent = is_persistent + self.page_size = page_size + self.is_var_seq = is_var_seq + self.is_var_split_kv = is_var_split_kv + self.cluster_shape_mnk = (2, 1, 1) + self.use_2cta_instrs = True + # When using 2 CTAs with m=128: warps 0-1 handle accumulation for first half [0, n/2), + # while warps 2-3 handle accumulation for second half [n/2, n) + self.warps_in_n = 2 + self.num_compute_warps = 4 + self.threads_per_warp = 32 + mma_qk_tiler_k = self.rope_dim + self.mma_qk_tiler = ( + self.mma_qk_tiler_mn[0], + self.mma_qk_tiler_mn[1], + mma_qk_tiler_k, + ) + self.mma_qk_rope_tiler = ( + self.mma_qk_tiler_mn[0], + self.mma_qk_tiler_mn[1], + self.rope_dim, + ) + self.mma_pv_tiler = ( + self.mma_pv_tiler_mn[0], + self.mma_pv_tiler_mn[1], + self.mma_qk_tiler[1] * self.mma_qk_tiler[2] // self.mma_pv_tiler_mn[1], + ) + self.iterations_qk_latent = self.latent_dim // self.mma_qk_tiler[2] + self.iterations_qk_rope = mma_qk_tiler_k // self.mma_qk_tiler[2] + self.iterations_qk = self.iterations_qk_latent + self.iterations_qk_rope + self.iterations_pv_k = self.mma_qk_tiler[1] // self.mma_pv_tiler[2] + self.iterations_pv_n = self.latent_dim // self.mma_pv_tiler[1] + + # Set specialized warp ids + self.compute_warp_ids = (0, 1, 2, 3) + self.correction_warp_ids = (4, 5, 6, 7) + self.mma_warp_id = 8 + + self.load_tma_warp_id = 9 + self.load_pt_warp_id = 10 + self.empty_warp_ids = (11,) + self.threads_per_cta = self.threads_per_warp * len( + ( + self.mma_warp_id, + self.load_tma_warp_id, + self.load_pt_warp_id, + *self.compute_warp_ids, + *self.correction_warp_ids, + *self.empty_warp_ids, + ) + ) + + # register settings + self.softmax_reg_num = 192 + self.correction_reg_num = 208 + self.other_reg_num = 96 + # Named barriers + self.tmem_ptr_sync_bar = pipeline.NamedBarrier( + barrier_id=1, + num_threads=( + self.threads_per_warp + + self.threads_per_warp * self.num_compute_warps * 2 + ), + ) + self.softmax_exchange_sync_bar = pipeline.NamedBarrier( + barrier_id=2, num_threads=(self.threads_per_warp * self.num_compute_warps) + ) + self.epilogue_exchange_sync_bar = pipeline.NamedBarrier( + barrier_id=3, num_threads=(self.threads_per_warp * self.num_compute_warps) + ) + + def _setup_attributes(self): + """Set up configurations and parameters for the MLA kernel operation. + + This method initializes and configures various attributes required for the + execution of the multi-head latent attention kernel, mainly about the pipeline stages: + + - Sets up staging parameters for Q, K, V inputs and accumulator data + - Configures pipeline stages for softmax, correction, and epilogue operations + """ + + self.load_q_stage = 1 + self.load_kv_stage = 15 + self.mma_s_stage = 2 + self.p_mma_stage = 2 + self.p_cor_stage = 2 + self.mma_o_stage = 1 + self.load_pt_stage = 4 + + self.tmem_o_offset = self.mma_s_stage * self.mma_qk_tiler[1] // self.warps_in_n + self.correction_factor_offset = ( + self.tmem_o_offset + self.latent_dim // self.warps_in_n + ) + + @cute.jit + def __call__( + self, + q_latent: cute.Tensor, + q_rope: cute.Tensor, + c_latent: cute.Tensor, + c_rope: cute.Tensor, + page_table: cute.Tensor, + o: cute.Tensor, + lse: cute.Tensor, + workspace: cute.Tensor, + split_kv: cutlass.Int32, + cache_seqs: Optional[cute.Tensor], + block_split_kvs: Optional[cute.Tensor], + softmax_scale: cutlass.Float32, + output_scale: cutlass.Float32, + stream: cuda.CUstream, + ): + """Execute the Multi-Head Latent Attention operation on the provided tensors. + + The method handles: + 1. Initialization of workspace for temporary split KV buffers + 2. Validation of tensor data types + 3. Initialization of hardware-specific parameters and memory layouts + 4. Configuration of TMA (Tensor Memory Access) operations + 5. Grid and work scheduling computation + 6. Kernel launch(split KV kernel and reduction kernel) with appropriate parameters + + :param q_latent: The query tensor with shape [num_head, latent_dim, seq_len_q, batch_size] + :type q_latent: cute.Tensor + :param q_rope: The query RoPE tensor with shape [num_head, rope_dim, seq_len_q, batch_size] + :type q_rope: cute.Tensor + :param c_latent: The key tensor with shape [seq_len_k, latent_dim, batch_size] + :type c_latent: cute.Tensor + :param c_rope: The key RoPE tensor with shape [seq_len_k, rope_dim, batch_size] + :type c_rope: cute.Tensor + :param page_table: The page table tensor with shape [page_count, batch_size] + :type page_table: cute.Tensor + :param o: The output tensor with shape [num_head, latent_dim, seq_len_q, batch_size] + :type o: cute.Tensor + :param lse: The LSE tensor with shape [num_head, seq_len_q, batch_size] + :type lse: cute.Tensor + :param workspace: The workspace tensor with 1-d shape prepared for acc_o and acc_lse + :type workspace: cute.Tensor + :param split_kv: The scalar factor for split KV + :type split_kv: cutlass.Int32 + :param cache_seqs: The cache sequences tensor with shape [batch_size] + :type cache_seqs: cute.Tensor + :param block_split_kvs: The block split KV tensor with shape [batch_size] + :type block_split_kvs: cute.Tensor + :param softmax_scale: The scale factor for softmax + :type softmax_scale: cutlass.Float32 + :param output_scale: The scale factor for the output + :type output_scale: cutlass.Float32 + :param stream: The CUDA stream to execute the kernel on + :type stream: cuda.CUstream + + :raises TypeError: If tensor data types don't match or aren't supported + """ + + # setup static attributes before smem/grid/tma computation + self.q_dtype = q_latent.element_type + self.k_dtype = c_latent.element_type + self.v_dtype = c_latent.element_type + self.o_dtype = o.element_type + + # check type consistency + if cutlass.const_expr( + self.q_dtype != self.k_dtype or self.q_dtype != self.v_dtype + ): + raise TypeError( + f"Type mismatch: {self.q_dtype} != {self.k_dtype} or {self.q_dtype} != {self.v_dtype}" + ) + # check leading dimensions of input/output + if cutlass.const_expr(q_latent.stride[1] != 1 or q_rope.stride[1] != 1): + raise ValueError("q_latent and q_rope must have leading dimension 1") + if cutlass.const_expr(c_latent.stride[1] != 1 or c_rope.stride[1] != 1): + raise ValueError("c_latent and c_rope must have leading dimension 1") + if cutlass.const_expr(o.stride[1] != 1): + raise ValueError("o must have leading dimension 1") + if cutlass.const_expr(lse.stride[0] != 1): + raise ValueError("lse must have leading dimension 0") + + acc_o, acc_lse = self.initialize_workspace( + q_latent.shape[0], + q_latent.shape[1], + q_latent.shape[2], + q_latent.shape[3], + split_kv, + self.acc_dtype, + workspace, + ) + + c_latent_tranpose_layout = cute.select(c_latent.layout, mode=[1, 0, 2]) + c_latent_transpose = cute.make_tensor( + c_latent.iterator, c_latent_tranpose_layout + ) + + self.q_major_mode = tcgen05.OperandMajorMode.K + self.k_major_mode = tcgen05.OperandMajorMode.K + self.v_major_mode = tcgen05.OperandMajorMode.MN + + self._setup_attributes() + + cta_group = tcgen05.CtaGroup.TWO + # the intermediate tensor p is from smem & k-major + p_major_mode = tcgen05.OperandMajorMode.K + qk_tiled_mma = sm100_utils.make_trivial_tiled_mma( + self.q_dtype, + self.q_major_mode, + self.k_major_mode, + self.acc_dtype, + cta_group, + self.mma_qk_tiler[:2], + ) + pv_tiled_mma = sm100_utils.make_trivial_tiled_mma( + self.v_dtype, + p_major_mode, + self.v_major_mode, + self.acc_dtype, + cta_group, + self.mma_pv_tiler[:2], + ) + + cta_layout_vmnk = cute.tiled_divide( + cute.make_layout(self.cluster_shape_mnk), + (qk_tiled_mma.thr_id.shape,), + ) + + self.epi_tile = self.mma_pv_tiler[:2] + + q_latent_smem_layout_staged = sm100_utils.make_smem_layout_a( + qk_tiled_mma, + self.mma_qk_tiler, + self.q_dtype, + (self.iterations_qk_latent * self.load_q_stage), + ) + q_latent_smem_layout_staged = cute.logical_divide( + q_latent_smem_layout_staged, (None, None, None, self.iterations_qk_latent) + ) + q_rope_smem_layout_staged = sm100_utils.make_smem_layout_a( + qk_tiled_mma, + self.mma_qk_rope_tiler, + self.q_dtype, + self.load_q_stage, + ) + + # rope reuse the same smem layout as latent + kc_smem_layout_staged = sm100_utils.make_smem_layout_b( + qk_tiled_mma, + self.mma_qk_tiler, + self.k_dtype, + self.load_kv_stage, + ) + kc_page_tile_size = min( + self.page_size, qk_tiled_mma.op.shape_mnk[0] // qk_tiled_mma.thr_id.shape + ) + + kc_smem_layout_for_tma = sm100_utils.make_smem_layout( + OperandMajorMode.K, + (self.mma_qk_tiler[0] // qk_tiled_mma.thr_id.shape, self.mma_qk_tiler[2]), + self.k_dtype, + self.load_kv_stage, + ) + kc_smem_layout_for_tma = cute.tiled_divide( + kc_smem_layout_for_tma, (kc_page_tile_size, self.mma_qk_tiler[2]) + ) + + p_smem_layout_staged = sm100_utils.make_smem_layout_a( + pv_tiled_mma, + self.mma_pv_tiler, + self.q_dtype, + (self.iterations_pv_k * self.p_mma_stage), + ) + p_smem_layout_staged = cute.logical_divide( + p_smem_layout_staged, (None, None, None, self.iterations_pv_k) + ) + + vc_smem_layout_staged = sm100_utils.make_smem_layout_b( + pv_tiled_mma, + self.mma_pv_tiler, + self.v_dtype, + self.load_kv_stage, + ) + vc_page_tile_size = min(self.page_size, self.mma_pv_tiler[2]) + vc_smem_layout_for_tma = sm100_utils.make_smem_layout( + OperandMajorMode.MN, + (self.mma_pv_tiler[1] // pv_tiled_mma.thr_id.shape, self.mma_pv_tiler[2]), + self.v_dtype, + self.load_kv_stage, + ) + vc_smem_layout_for_tma = cute.tiled_divide( + vc_smem_layout_for_tma, + ( + pv_tiled_mma.op.shape_mnk[1] // pv_tiled_mma.thr_id.shape, + vc_page_tile_size, + ), + ) + # TMA load for Q latent and rope + tma_load_op = cute.nvgpu.cpasync.CopyBulkTensorTileG2SOp(cta_group) + + q_latent_smem_layout = cute.select(q_latent_smem_layout_staged, mode=[0, 1, 2]) + tma_atom_q_latent, tma_tensor_q_latent = cute.nvgpu.make_tiled_tma_atom_A( + tma_load_op, + q_latent, + q_latent_smem_layout, + self.mma_qk_tiler, + qk_tiled_mma, + cta_layout_vmnk.shape, + ) + q_rope_smem_layout = cute.select(q_rope_smem_layout_staged, mode=[0, 1, 2]) + tma_atom_q_rope, tma_tensor_q_rope = cute.nvgpu.make_tiled_tma_atom_A( + tma_load_op, + q_rope, + q_rope_smem_layout, + self.mma_qk_rope_tiler, + qk_tiled_mma, + cta_layout_vmnk.shape, + ) + # TMA load for c latent and k rope + kc_smem_layout = cute.select(kc_smem_layout_for_tma, mode=[0]) + tma_atom_c_latent, tma_tensor_c_latent = self.make_paged_tiled_tma_atom( + tma_load_op, + c_latent, + kc_smem_layout, + (self.mma_qk_tiler[1], self.mma_qk_tiler[2]), + qk_tiled_mma, + is_k_load=True, + ) + tma_atom_c_rope, tma_tensor_c_rope = self.make_paged_tiled_tma_atom( + tma_load_op, + c_rope, + kc_smem_layout, + (self.mma_qk_tiler[1], self.mma_qk_tiler[2]), + qk_tiled_mma, + is_k_load=True, + ) + # TMA load for c latent transpose + vc_smem_layout = cute.select(vc_smem_layout_for_tma, mode=[0]) + tma_atom_c_latent_transpose, tma_tensor_c_latent_transpose = ( + self.make_paged_tiled_tma_atom( + tma_load_op, + c_latent_transpose, + vc_smem_layout, + (self.mma_pv_tiler[1], self.mma_pv_tiler[2]), + pv_tiled_mma, + is_k_load=False, + ) + ) + + q_latent_copy_size = ( + cute.size_in_bytes(self.q_dtype, q_latent_smem_layout) + * cute.size(qk_tiled_mma.thr_id.shape) + * self.iterations_qk_latent + ) + q_rope_copy_size = ( + cute.size_in_bytes(self.q_dtype, q_rope_smem_layout) + * cute.size(qk_tiled_mma.thr_id.shape) + * self.iterations_qk_rope + ) + q_copy_size = q_latent_copy_size + q_rope_copy_size + kc_copy_size = cute.size_in_bytes( + self.k_dtype, cute.select(kc_smem_layout_staged, mode=[0, 1, 2]) + ) * cute.size(qk_tiled_mma.thr_id.shape) + vc_copy_size = cute.size_in_bytes( + self.v_dtype, cute.select(vc_smem_layout_staged, mode=[0, 1, 2]) + ) * cute.size(pv_tiled_mma.thr_id.shape) + assert ( + kc_copy_size == vc_copy_size + ), "kc_copy_size and vc_copy_size must be the same" + + self.tma_copy_q_bytes = q_copy_size + self.tma_copy_kc_bytes = kc_copy_size + + tile_sched_params, grid = self._compute_grid( + o, + split_kv, + self.cluster_shape_mnk, + self.max_active_clusters, + self.is_persistent, + ) + + @cute.struct + class SplitKVKernelSharedStorage: + # Pipeline barriers + load_q_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.load_q_stage * 2] + load_kv_mbar_ptr: cute.struct.MemRange[ + cutlass.Int64, self.load_kv_stage * 2 + ] + mma_s_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.mma_s_stage * 2] + p_mma_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.p_mma_stage * 2] + p_cor_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.p_cor_stage * 2] + mma_o_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.mma_o_stage * 2] + load_pt_mbar_ptr: cute.struct.MemRange[ + cutlass.Int64, self.load_pt_stage * 2 + ] + # Tmem dealloc cluster barrier + tmem_dealloc_mbar_ptr: cutlass.Int64 + + # Tmem holding buffer + tmem_holding_buf: cutlass.Int32 + # Smem tensors + softmax_smem_exchange: cute.struct.MemRange[ + self.acc_dtype, self.num_compute_warps * self.threads_per_warp + ] + epilogue_smem_exchange: cute.struct.MemRange[ + self.acc_dtype, self.num_compute_warps * self.threads_per_warp + ] + smem_q_latent: cute.struct.Align[ + cute.struct.MemRange[ + self.q_dtype, cute.cosize(q_latent_smem_layout_staged) + ], + 1024, + ] + smem_q_rope: cute.struct.Align[ + cute.struct.MemRange[ + self.q_dtype, cute.cosize(q_rope_smem_layout_staged) + ], + 1024, + ] + smem_kc: cute.struct.Align[ + cute.struct.MemRange[self.k_dtype, cute.cosize(kc_smem_layout_staged)], + 1024, + ] + smem_p: cute.struct.Align[ + cute.struct.MemRange[self.q_dtype, cute.cosize(p_smem_layout_staged)], + 1024, + ] + smem_page_table: cute.struct.MemRange[ + cutlass.Int32, self.load_pt_stage * self.mma_qk_tiler[1] // 2 + ] + + softmax_scale_log2 = softmax_scale * LOG2_E + self.split_kv_kernel( + qk_tiled_mma, + pv_tiled_mma, + tma_atom_q_latent, + tma_tensor_q_latent, + tma_atom_q_rope, + tma_tensor_q_rope, + tma_atom_c_latent, + tma_tensor_c_latent, + tma_atom_c_rope, + tma_tensor_c_rope, + tma_atom_c_latent_transpose, + tma_tensor_c_latent_transpose, + page_table, + o, + lse, + acc_o, + acc_lse, + split_kv, + cache_seqs, + block_split_kvs, + softmax_scale_log2, + output_scale, + q_latent_smem_layout_staged, + q_rope_smem_layout_staged, + kc_smem_layout_staged, + p_smem_layout_staged, + vc_smem_layout_staged, + kc_smem_layout_for_tma, + vc_smem_layout_for_tma, + cta_layout_vmnk, + tile_sched_params, + SplitKVKernelSharedStorage, + ).launch( + grid=grid, + block=[self.threads_per_cta, 1, 1], + cluster=self.cluster_shape_mnk, + smem=SplitKVKernelSharedStorage.size_in_bytes(), + stream=stream, + min_blocks_per_mp=1, + ) + if cutlass.const_expr(acc_o is not None): + self.reduction_kernel( + o, + lse, + acc_o, + acc_lse, + split_kv, + cache_seqs, + block_split_kvs, + ).launch( + grid=(q_latent.shape[0], q_latent.shape[2], q_latent.shape[3]), + block=[self.threads_per_warp * self.num_compute_warps, 1, 1], + smem=MAX_SPLITS * self.acc_dtype.width // 8, + stream=stream, + min_blocks_per_mp=1, + ) + + @cute.jit + def make_paged_tiled_tma_atom( + self, + tma_load_op: cute.nvgpu.cpasync.CopyBulkTensorTileG2SOp, + gmem: cute.Tensor, + smem_layout: cute.Layout, + mma_tiler, + tiled_mma: cute.TiledMma, + is_k_load: bool, + ): + ident = cute.make_identity_layout(gmem.shape) + g_tile = cute.composition(ident, mma_tiler) + cta_mn = mma_tiler[0] // tiled_mma.thr_id.shape + cta_v_map = cute.flat_divide(g_tile, (cta_mn,)) + cta_v_map = cute.select(cta_v_map, mode=[0, 2]) + page_tile_size = ( + min(self.page_size, cta_mn) + if is_k_load + else min(self.page_size, mma_tiler[1]) + ) + cta_v_map = cute.zipped_divide( + cta_v_map, + (page_tile_size, mma_tiler[1]) if is_k_load else (cta_mn, page_tile_size), + ) + cta_v_map = cute.select(cta_v_map, mode=[0]) + from cutlass._mlir.dialects import cute_nvgpu as _cute_nvgpu_ir + + res = _cute_nvgpu_ir.atom_make_non_exec_tiled_tma_load( + gmem.value, + smem_layout.value, + cta_v_map, + tma_load_op._to_ir(), + num_multicast=1, + ) + return ( + cute.CopyAtom( + tma_load_op, cpasync.CopyBulkTensorTileG2SNonExecTrait(res[0]) + ), + res[1], + ) + + @cute.kernel + def split_kv_kernel( + self, + tiled_mma_qk: cute.TiledMma, + tiled_mma_pv: cute.TiledMma, + tma_atom_q_latent: Optional[cute.CopyAtom], + mQL: cute.Tensor, + tma_atom_q_rope: Optional[cute.CopyAtom], + mQR: cute.Tensor, + tma_atom_c_latent: Optional[cute.CopyAtom], + mCL: cute.Tensor, + tma_atom_c_rope: Optional[cute.CopyAtom], + mKR: cute.Tensor, + tma_atom_c_latent_transpose: Optional[cute.CopyAtom], + mCLT: cute.Tensor, + mPT: cute.Tensor, + mO: Optional[cute.Tensor], + mLSE: Optional[cute.Tensor], + mAccO: Optional[cute.Tensor], + mAccLSE: Optional[cute.Tensor], + split_kv: cutlass.Int32, + cache_seqs: cute.Tensor, + block_split_kvs: cute.Tensor, + softmax_scale_log2: cutlass.Float32, + output_scale: cutlass.Float32, + q_latent_smem_layout_staged: cute.ComposedLayout, + q_rope_smem_layout_staged: cute.ComposedLayout, + kc_smem_layout_staged: cute.ComposedLayout, + p_smem_layout_staged: cute.ComposedLayout, + vc_smem_layout_staged: cute.ComposedLayout, + kc_smem_layout_for_tma: cute.ComposedLayout, + vc_smem_layout_for_tma: cute.ComposedLayout, + cta_layout_vmnk: cute.Layout, + tile_sched_params: MLAStaticTileSchedulerParams, + SharedStorage: cutlass.Constexpr, + ): + """The device split_kv kernel implementation of the Multi-Head Latent Attention. + + This kernel coordinates multiple specialized warps to perform different phases of the MLA computation: + 1. Load warp: Loads Q/C latent/rope data from global memory to shared memory using TMA + 2. MMA warp: Performs matrix multiplications (Q*K^T and P*V) + 3. Compute warps: Compute softmax and do rescaling on accumulators, and store the intermediate/final results + to global memory + + The kernel produces either intermediate or final results of the MLA computation based on the split_kv parameter. + When split_kv is 1, the kernel generates the final results directly. Otherwise, it produces intermediate results + that will later be combined by a reduction kernel. + + The kernel implements a complex pipeline with overlapping computation and memory operations, + using tensor memory access (TMA) for efficient data loading, warp specialization for different + computation phases. + + :param tiled_mma_qk: Tiled MMA for Q*K^T + :type tiled_mma_qk: cute.TiledMma + :param tiled_mma_pv: Tiled MMA for P*V + :type tiled_mma_pv: cute.TiledMma + :param tma_atom_q_latent: TMA copy atom for query latent tensor + :type tma_atom_q_latent: cute.CopyAtom + :param mQL: query latent tensor + :type mQL: cute.Tensor + :param tma_atom_q_rope: TMA copy atom for query rope tensor + :type tma_atom_q_rope: cute.CopyAtom + :param mKR: Compressed rope tensor + :type mKR: cute.Tensor + :param tma_atom_c_latent: TMA copy atom for c latent tensor + :type tma_atom_c_latent: cute.CopyAtom + :param mCL: Compressed latent tensor + :type mCL: cute.Tensor + :param tma_atom_c_rope: TMA copy atom for c rope tensor + :type tma_atom_c_rope: cute.CopyAtom + :param mCLT: Compressed latent transpose tensor + :type mCLT: cute.Tensor + :param mPT: Page table tensor + :type mPT: cute.Tensor + :param mO: Output tensor + :type mO: cute.Tensor + :param mLSE: Log-sum-exp tensor + :type mLSE: cute.Tensor + :param mAccO: Intermediate accumulator output tensor + :type mAccO: cute.Tensor + :param mAccLSE: Intermediate accumulator log-sum-exp tensor + :type mAccLSE: cute.Tensor + :param split_kv: The split_kv parameter + :type split_kv: cutlass.Int32 + :param cache_seqs: The variable sequence length tensor + :type cache_seqs: cute.Tensor + :param block_split_kvs: The per-block split_kv values tensor + :type block_split_kvs: cute.Tensor + :param softmax_scale_log2: The log2 scale factor for softmax + :type softmax_scale_log2: cutlass.Float32 + :param output_scale: The scale factor for the output + :type output_scale: cutlass.Float32 + :param q_latent_smem_layout_staged: Shared memory layout for query latent tensor + :type q_latent_smem_layout_staged: cute.ComposedLayout + :param q_rope_smem_layout_staged: Shared memory layout for query rope tensor + :type q_rope_smem_layout_staged: cute.ComposedLayout + :param kc_smem_layout_staged: Shared memory layout for key/value latent/rope tensor + :type kc_smem_layout_staged: cute.ComposedLayout + :param p_smem_layout_staged: Shared memory layout for probability matrix + :type p_smem_layout_staged: cute.ComposedLayout + :param vc_smem_layout_staged: Shared memory layout for value tensor + :type vc_smem_layout_staged: cute.ComposedLayout + :param kc_smem_layout_for_tma: Shared memory layout for key/value latent tensor for TMA + :type kc_smem_layout_for_tma: cute.ComposedLayout + :param vc_smem_layout_for_tma: Shared memory layout for value tensor for TMA + :type vc_smem_layout_for_tma: cute.ComposedLayout + :param cta_layout_vmnk: Layout for compute threads + :type cta_layout_vmnk: cute.Layout + :param tile_sched_params: Scheduling parameters for work distribution + :type tile_sched_params: MLAStaticTileSchedulerParams + :param SharedStorage: Shared storage for the kernel + :type SharedStorage: cutlass.Constexpr + """ + + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + + tidx, _, _ = cute.arch.thread_idx() + bidx, _, _ = cute.arch.block_idx() + mma_tile_coord_v = bidx % cute.size(tiled_mma_qk.thr_id.shape) + is_leader_cta = mma_tile_coord_v == 0 + + # Prefetch tma descriptor + if warp_idx == self.mma_warp_id: + cpasync.prefetch_descriptor(tma_atom_q_latent) + cpasync.prefetch_descriptor(tma_atom_q_rope) + cpasync.prefetch_descriptor(tma_atom_c_latent) + cpasync.prefetch_descriptor(tma_atom_c_rope) + cpasync.prefetch_descriptor(tma_atom_c_latent_transpose) + + # Alloc + smem = utils.SmemAllocator() + storage = smem.allocate(SharedStorage) + + # Tensor memory dealloc barrier init + tmem = utils.TmemAllocator( + storage.tmem_holding_buf, + barrier_for_retrieve=self.tmem_ptr_sync_bar, + allocator_warp_id=self.mma_warp_id, + is_two_cta=self.use_2cta_instrs, + two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar_ptr, + ) + + load_q_pipeline = self.make_and_init_load_qkv_pipeline( + storage.load_q_mbar_ptr.data_ptr(), + cta_layout_vmnk, + self.load_q_stage, + self.tma_copy_q_bytes, + ) + load_kv_pipeline = self.make_and_init_load_qkv_pipeline( + storage.load_kv_mbar_ptr.data_ptr(), + cta_layout_vmnk, + self.load_kv_stage, + self.tma_copy_kc_bytes, + ) + mma_s_pipeline = self.make_and_init_mma_s_pipeline( + storage.mma_s_mbar_ptr.data_ptr(), cta_layout_vmnk + ) + p_mma_pipeline = self.make_and_init_p_mma_pipeline( + storage.p_mma_mbar_ptr.data_ptr(), cta_layout_vmnk + ) + p_cor_pipeline = self.make_and_init_p_cor_pipeline( + storage.p_cor_mbar_ptr.data_ptr() + ) + mma_o_pipeline = self.make_and_init_mma_o_pipeline( + storage.mma_o_mbar_ptr.data_ptr(), cta_layout_vmnk + ) + load_pt_pipeline = self.make_and_init_load_pt_pipeline( + storage.load_pt_mbar_ptr.data_ptr() + ) + + # Cluster arrive after barrier init + pipeline_init_arrive(cluster_shape_mn=self.cluster_shape_mnk, is_relaxed=True) + + # Generate smem tensor Q/KC/VC/exchange + # (MMA, MMA_H, MMA_R, PIPE) + sQ = storage.smem_q_latent.get_tensor( + q_latent_smem_layout_staged.outer, swizzle=q_latent_smem_layout_staged.inner + ) + sQ_rope = storage.smem_q_rope.get_tensor( + q_rope_smem_layout_staged.outer, swizzle=q_rope_smem_layout_staged.inner + ) + # (MMA, MMA_K, MMA_R, PIPE) + sKC = storage.smem_kc.get_tensor( + kc_smem_layout_staged.outer, swizzle=kc_smem_layout_staged.inner + ) + sKC_for_tma = storage.smem_kc.get_tensor( + kc_smem_layout_for_tma.outer, + swizzle=kc_smem_layout_for_tma.inner, + ) + # (MMA, MMA_D, MMA_K, PIPE) + # reuse smem + sVC_ptr = cute.recast_ptr(sKC.iterator, vc_smem_layout_staged.inner) + sVC = cute.make_tensor(sVC_ptr, vc_smem_layout_staged.outer) + sVC_for_tma = cute.make_tensor(sVC_ptr, vc_smem_layout_for_tma.outer) + # (MMA, MMA_H, MMA_K) + sP = storage.smem_p.get_tensor( + p_smem_layout_staged.outer, swizzle=p_smem_layout_staged.inner + ) + sPT = storage.smem_page_table.get_tensor( + cute.make_layout((self.mma_qk_tiler[1] // 2, self.load_pt_stage)) + ) + # (compute_threads,) + softmax_smem_exchange = storage.softmax_smem_exchange.get_tensor( + cute.make_layout(self.num_compute_warps * self.threads_per_warp) + ) + epilogue_smem_exchange = storage.epilogue_smem_exchange.get_tensor( + cute.make_layout(self.num_compute_warps * self.threads_per_warp) + ) + + # + # Cluster wait before tensor memory alloc + # + pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mnk) + + # /////////////////////////////////////////////////////////////////////////////// + # Load warps, including page table and data tensors + # /////////////////////////////////////////////////////////////////////////////// + + if warp_idx >= self.empty_warp_ids[0] and warp_idx <= self.empty_warp_ids[-1]: + cute.arch.setmaxregister_decrease(self.other_reg_num) + if warp_idx == self.load_pt_warp_id: + cute.arch.setmaxregister_decrease(self.other_reg_num) + load_pt_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.load_pt_stage + ) + tile_sched = create_mla_static_tile_scheduler( + tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim() + ) + work_tile = tile_sched.initial_work_tile_info() + while work_tile.is_valid_tile: + blk_coord = work_tile.tile_idx + k_index, k_tile_count, local_split_kv = self.get_k_tile_count( + split_kv, + cache_seqs, + block_split_kvs, + blk_coord, + ) + if k_tile_count > 0: + load_pt_common_params = SimpleNamespace( + blk_coord=blk_coord, + load_pt_pipeline=load_pt_pipeline, + mPT=mPT, + sPT=sPT, + tidx=tidx, + page_size=mCL.shape[0], + ) + load_pt_producer_state = self.load_page_table( + load_pt_common_params, + k_index, + k_tile_count, + load_pt_producer_state, + ) + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + load_pt_pipeline.producer_tail(load_pt_producer_state) + if warp_idx == self.load_tma_warp_id: + cute.arch.setmaxregister_decrease(self.other_reg_num) + load_q_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.load_q_stage + ) + load_kv_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.load_kv_stage + ) + load_pt_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.load_pt_stage + ) + load_pt_release_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.load_pt_stage + ) + tile_sched = create_mla_static_tile_scheduler( + tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim() + ) + work_tile = tile_sched.initial_work_tile_info() + while work_tile.is_valid_tile: + blk_coord = work_tile.tile_idx + k_index, k_tile_count, local_split_kv = self.get_k_tile_count( + split_kv, + cache_seqs, + block_split_kvs, + blk_coord, + ) + if k_tile_count > 0: + # Construct fixed common/tma_qk/tma_pv params for load_tma + tma_common_params = SimpleNamespace( + blk_coord=blk_coord, + local_split_kv=local_split_kv, + load_q_pipeline=load_q_pipeline, + load_kv_pipeline=load_kv_pipeline, + mPT=mPT, + sPT=sPT, + load_pt_pipeline=load_pt_pipeline, + ) + tma_qk_params = SimpleNamespace( + tiled_mma_qk=tiled_mma_qk, + tma_atom_q_latent=tma_atom_q_latent, + tma_atom_q_rope=tma_atom_q_rope, + tma_atom_c_latent=tma_atom_c_latent, + tma_atom_c_rope=tma_atom_c_rope, + mQL=mQL, + mQR=mQR, + mCL=mCL, + mKR=mKR, + sQ=sQ, + sQ_rope=sQ_rope, + sKC=sKC_for_tma, + ) + tma_pv_params = SimpleNamespace( + tiled_mma_pv=tiled_mma_pv, + tma_atom_c_latent_transpose=tma_atom_c_latent_transpose, + mCL=mCL, + mKR=mKR, + mCLT=mCLT, + sVC=sVC_for_tma, + ) + # Load tma + ( + load_q_producer_state, + load_kv_producer_state, + load_pt_consumer_state, + load_pt_release_state, + ) = self.load_tma( + tma_common_params, + tma_qk_params, + tma_pv_params, + k_index, + k_tile_count, + load_q_producer_state, + load_kv_producer_state, + load_pt_consumer_state, + load_pt_release_state, + ) + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + + load_q_pipeline.producer_tail(load_q_producer_state) + load_kv_pipeline.producer_tail(load_kv_producer_state) + + # /////////////////////////////////////////////////////////////////////////////// + # MMA warp + # /////////////////////////////////////////////////////////////////////////////// + if warp_idx == self.mma_warp_id: + cute.arch.setmaxregister_decrease(self.other_reg_num) + # Alloc tensor memory buffer + tmem.allocate(cute.arch.get_max_tmem_alloc_cols("sm_100")) + tmem.wait_for_alloc() + tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + + load_q_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.load_q_stage + ) + load_kv_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.load_kv_stage + ) + mma_s_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.mma_s_stage + ) + p_mma_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.p_mma_stage + ) + mma_o_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.mma_o_stage + ) + tile_sched = create_mla_static_tile_scheduler( + tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim() + ) + work_tile = tile_sched.initial_work_tile_info() + while work_tile.is_valid_tile: + blk_coord = work_tile.tile_idx + k_index, k_tile_count, local_split_kv = self.get_k_tile_count( + split_kv, cache_seqs, block_split_kvs, blk_coord + ) + if k_tile_count > 0: + mma_common_params = SimpleNamespace( + blk_coord=blk_coord, + local_split_kv=local_split_kv, + load_q_pipeline=load_q_pipeline, + load_kv_pipeline=load_kv_pipeline, + tmem_ptr=tmem_ptr, + is_leader_cta=is_leader_cta, + L=mCL.shape[1], + ) + mma_qk_params = SimpleNamespace( + mma_s_pipeline=mma_s_pipeline, + sQ=sQ, + sQ_rope=sQ_rope, + sKC=sKC, + ) + mma_pv_params = SimpleNamespace( + p_mma_pipeline=p_mma_pipeline, + mma_o_pipeline=mma_o_pipeline, + sP=sP, + sVC=sVC, + ) + ( + tiled_mma_qk, + tiled_mma_pv, + load_q_consumer_state, + load_kv_consumer_state, + mma_s_producer_state, + p_mma_consumer_state, + mma_o_producer_state, + ) = self.mma( + mma_common_params, + mma_qk_params, + mma_pv_params, + k_tile_count, + tiled_mma_qk, + tiled_mma_pv, + load_q_consumer_state, + load_kv_consumer_state, + mma_s_producer_state, + p_mma_consumer_state, + mma_o_producer_state, + ) + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + + mma_s_pipeline.producer_tail(mma_s_producer_state) + mma_o_pipeline.producer_tail(mma_o_producer_state) + + tmem.relinquish_alloc_permit() + tmem.free(tmem_ptr) + + # /////////////////////////////////////////////////////////////////////////////// + # Compute warp + # /////////////////////////////////////////////////////////////////////////////// + if ( + warp_idx >= self.compute_warp_ids[0] + and warp_idx <= self.compute_warp_ids[-1] + ): + cute.arch.setmaxregister_increase(self.softmax_reg_num) + mma_s_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.mma_s_stage + ) + p_mma_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.p_mma_stage + ) + p_cor_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.p_cor_stage + ) + mma_o_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.mma_o_stage + ) + # sync with mma warp before retrieving tmem ptr + tmem.wait_for_alloc() + + tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + + tile_sched = create_mla_static_tile_scheduler( + tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim() + ) + work_tile = tile_sched.initial_work_tile_info() + while work_tile.is_valid_tile: + blk_coord = work_tile.tile_idx + k_index, k_tile_count, local_split_kv = self.get_k_tile_count( + split_kv, cache_seqs, block_split_kvs, blk_coord + ) + if k_tile_count > 0: + compute_common_params = SimpleNamespace( + blk_coord=blk_coord, + split_kv=split_kv, + local_split_kv=local_split_kv, + smem_exchange=softmax_smem_exchange, + mAccO=mAccO, + mO=mO, + K=cache_seqs[blk_coord[2]], + L=mCL.shape[1], + tmem_ptr=tmem_ptr, + tidx=tidx, + p_cor_pipeline=p_cor_pipeline, + ) + compute_softmax_params = SimpleNamespace( + tiled_mma_qk=tiled_mma_qk, + sP=sP, + mma_s_pipeline=mma_s_pipeline, + p_mma_pipeline=p_mma_pipeline, + softmax_scale_log2=softmax_scale_log2, + ) + mma_s_consumer_state, p_mma_producer_state, p_cor_producer_state = ( + self.compute( + compute_common_params, + compute_softmax_params, + k_index=k_index, + k_tile_count=k_tile_count, + mma_s_consumer_state=mma_s_consumer_state, + p_mma_producer_state=p_mma_producer_state, + p_cor_producer_state=p_cor_producer_state, + ) + ) + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + p_cor_pipeline.producer_tail(p_cor_producer_state) + + # /////////////////////////////////////////////////////////////////////////////// + # Correction warp + # /////////////////////////////////////////////////////////////////////////////// + if ( + warp_idx >= self.correction_warp_ids[0] + and warp_idx <= self.correction_warp_ids[-1] + ): + cute.arch.setmaxregister_increase(self.correction_reg_num) + p_cor_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.p_cor_stage + ) + mma_o_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.mma_o_stage + ) + # sync with mma warp before retrieving tmem ptr + tmem.wait_for_alloc() + + tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + + tile_sched = create_mla_static_tile_scheduler( + tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim() + ) + work_tile = tile_sched.initial_work_tile_info() + while work_tile.is_valid_tile: + blk_coord = work_tile.tile_idx + k_index, k_tile_count, local_split_kv = self.get_k_tile_count( + split_kv, cache_seqs, block_split_kvs, blk_coord + ) + if k_tile_count > 0: + compute_common_params = SimpleNamespace( + blk_coord=blk_coord, + split_kv=split_kv, + local_split_kv=local_split_kv, + smem_exchange=epilogue_smem_exchange, + mAccO=mAccO, + mO=mO, + K=cache_seqs[blk_coord[2]], + L=mCL.shape[1], + H=mQL.shape[0], + tmem_ptr=tmem_ptr, + tidx=tidx, + tiled_mma_pv=tiled_mma_pv, + p_cor_pipeline=p_cor_pipeline, + mma_o_pipeline=mma_o_pipeline, + ) + compute_epilogue_params = SimpleNamespace( + output_scale=output_scale, + softmax_scale_log2=softmax_scale_log2, + mAccLSE=mAccLSE, + mLSE=mLSE, + ) + p_cor_consumer_state, mma_o_consumer_state = self.correction( + compute_common_params, + compute_epilogue_params, + k_tile_count=k_tile_count, + p_cor_consumer_state=p_cor_consumer_state, + mma_o_consumer_state=mma_o_consumer_state, + ) + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + + return + + @cute.kernel + def reduction_kernel( + self, + mO: cute.Tensor, + mLSE: cute.Tensor, + mAccO: cute.Tensor, + mAccLSE: cute.Tensor, + split_kv: cutlass.Int32, + cache_seqs: cute.Tensor, + block_split_kvs: cute.Tensor, + ): + """The reduction kernel for Multi-Head Latent Attention (MLA) that combines intermediate results + from multiple split_kv blocks into final outputs. + + :param mO: Output tensor for storing final results + :type mO: cute.Tensor + :param mLSE: Log-sum-exp tensor for storing final LSE values + :type mLSE: cute.Tensor + :param mAccO: Accumulated output tensor from split_kv blocks + :type mAccO: cute.Tensor + :param mAccLSE: Accumulated LSE tensor from split_kv blocks + :type mAccLSE: cute.Tensor + :param split_kv: Number of split_kv blocks + :type split_kv: cutlass.Int32 + :param cache_seqs: Cache sequence lengths tensor + :type cache_seqs: cute.Tensor + :param block_split_kvs: Per-block split_kv values tensor (for variable split_kv) + :type block_split_kvs: cute.Tensor + """ + bidx, bidy, bidz = cute.arch.block_idx() + tidx, _, _ = cute.arch.thread_idx() + blk_coord = (bidx, bidy, bidz) + local_split_kv = ( + block_split_kvs[blk_coord[2]] if self.is_var_split_kv else split_kv + ) + k_tile_total = cute.ceil_div(cache_seqs[blk_coord[2]], self.mma_qk_tiler[1]) + k_tile_per_cta = cute.ceil_div(k_tile_total, local_split_kv) + local_split_kv = cute.ceil_div(k_tile_total, k_tile_per_cta) + + # Alloc shared memory + smem = utils.SmemAllocator() + storage = smem.allocate(MAX_SPLITS * self.acc_dtype.width // 8, 16) + lse_scale_ptr = cute.recast_ptr(storage, dtype=self.acc_dtype) + smem_lse_scale = cute.make_tensor(lse_scale_ptr, cute.make_layout(MAX_SPLITS)) + + gLSE = mAccLSE[blk_coord[0], None, blk_coord[1], blk_coord[2]] + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + if warp_idx == 0: + # calculate the global lse and exp ^ (local_lse - global_lse) + lse_per_thread = cute.ceil_div(MAX_SPLITS, self.threads_per_warp) + + local_lse = cute.make_rmem_tensor( + cute.make_layout(lse_per_thread), self.lse_dtype + ) + lse_max = -self.lse_dtype.inf + # find the max lse + for i in cutlass.range_constexpr(lse_per_thread): + split_kv_idx = tidx + i * self.threads_per_warp + local_lse[i] = ( + gLSE[split_kv_idx] + if cute.elem_less(split_kv_idx, local_split_kv) + else -self.lse_dtype.inf + ) + # reduce the local lse + lse_max = cute.arch.fmax(lse_max, local_lse[i]) + lse_max = cute.arch.warp_reduction_max(lse_max) + lse_max = lse_max if lse_max != -self.lse_dtype.inf else 0.0 + # calculate sum_lse + sum_lse = 0.0 + for i in cutlass.range_constexpr(lse_per_thread): + sum_lse += cute.math.exp2(local_lse[i] - lse_max, fastmath=True) + sum_lse = cute.arch.warp_reduction_sum(sum_lse) + # calculate the global_lse + global_lse = ( + lse_max + cute.math.log2(sum_lse, fastmath=True) + if not sum_lse == self.lse_dtype(0.0) or sum_lse != sum_lse + else self.lse_dtype.inf + ) + if tidx == 0: + mLSE[blk_coord[0], blk_coord[1], blk_coord[2]] = global_lse + # store the scale to shared memory + for i in cutlass.range_constexpr(lse_per_thread): + split_kv_idx = tidx + i * self.threads_per_warp + if cute.elem_less(split_kv_idx, local_split_kv): + smem_lse_scale[split_kv_idx] = cute.math.exp2( + local_lse[i] - global_lse, fastmath=True + ) + + pipeline.sync(barrier_id=4) + + elements_per_thread = cute.ceil_div( + self.latent_dim, self.threads_per_warp * self.num_compute_warps + ) + gAccO = mAccO[blk_coord[0], None, None, blk_coord[1], blk_coord[2]] + rAccO = cute.make_rmem_tensor( + cute.make_layout(elements_per_thread), self.acc_dtype + ) + rO = cute.make_rmem_tensor(cute.make_layout(elements_per_thread), self.o_dtype) + rAccO.fill(0.0) + for i in range(local_split_kv): + for j in cutlass.range_constexpr(elements_per_thread): + element_idx = tidx + j * self.threads_per_warp * self.num_compute_warps + rAccO[j] += gAccO[i, element_idx] * smem_lse_scale[i] + rO.store(rAccO.load().to(self.o_dtype)) + for j in cutlass.range_constexpr(elements_per_thread): + element_idx = tidx + j * self.threads_per_warp * self.num_compute_warps + mO[blk_coord[0], element_idx, blk_coord[1], blk_coord[2]] = rO[j] + return + + @staticmethod + def get_split_kv( + B: int, S: int, K: int, mma_qk_tiler_mn: tuple, max_active_blocks: int + ) -> int: + """Get the proper split_kv value for the MLA kernel based on parameters. + + :param B: Batch size + :type B: int + :param S: Sequence length + :type S: int + :param K: Sequence length + :type K: int + :param mma_qk_tiler_mn: MLA tiling parameters + :type mma_qk_tiler_mn: tuple + :param max_active_blocks: Maximum number of active blocks + :type max_active_blocks: int + :return: Split_kv value + :rtype: int + """ + max_splits = ceil_div(K, mma_qk_tiler_mn[1]) + blocks_per_batch = max(1, max_active_blocks // B // (S * 2)) + split_heur = min(max_splits, blocks_per_batch) + k_waves = ceil_div(max_splits, split_heur) + split_wave_aware = ceil_div(max_splits, k_waves) + max_split_kv = 32 + return min(split_wave_aware, max_split_kv) + + @cute.jit + def get_k_tile_count( + self, + split_kv: cutlass.Int32, + cache_seqs: cute.Tensor, + block_split_kvs: cute.Tensor, + blk_coord: cute.Coord, + ) -> tuple[cutlass.Int32, cutlass.Int32, cutlass.Int32]: + """Get the current k_index, k_tile_count, and local split_kv value for the MLA kernel. + + :param split_kv: Split_kv value + :type split_kv: cutlass.Int32 + :param cache_seqs: Cache sequence lengths tensor + :type cache_seqs: cute.Tensor + :param block_split_kvs: Per-block split_kv values tensor + :type block_split_kvs: cute.Tensor + :param blk_coord: Block coordinate + :type blk_coord: cute.Coord + :return: k_index, k_tile_count, split_kv + :rtype: tuple[cutlass.Int32, cutlass.Int32, cutlass.Int32] + """ + K = cache_seqs[blk_coord[2]] + if cutlass.const_expr(self.is_var_split_kv): + split_kv = block_split_kvs[blk_coord[2]] + + k_tile_total = cute.ceil_div(K, self.mma_qk_tiler[1]) + k_tile_per_cta = cute.ceil_div(k_tile_total, split_kv) + k_index = blk_coord[3] * k_tile_per_cta + k_tile_count = max(0, min(k_tile_total, k_index + k_tile_per_cta) - k_index) + return k_index, k_tile_count, split_kv + + @cute.jit + def load_page_table( + self, + common_params: SimpleNamespace, + k_index: cutlass.Int32, + k_tile_count: cutlass.Int32, + load_pt_producer_state: pipeline.PipelineState, + ) -> pipeline.PipelineState: + """Load warp to load page table. Updates the load pt producer state. + + :param common_params: The common parameters + :type common_params: SimpleNamespace + :param k_index: The k index + :type k_index: cutlass.Int32 + :param k_tile_count: The k tile count + :type k_tile_count: cutlass.Int32 + :param load_pt_producer_state: The load pt producer state + :type load_pt_producer_state: pipeline.PipelineState + + :return: The load pt producer state + :rtype: pipeline.PipelineState + """ + mPT = common_params.mPT[None, common_params.blk_coord[2]] + page_per_tile = self.mma_qk_tiler[1] // self.page_size + tidx = common_params.tidx % self.threads_per_warp + + load_pt_pipeline = common_params.load_pt_pipeline + while k_tile_count > 0: + load_pt_pipeline.producer_acquire(load_pt_producer_state) + + elem_per_thread = cute.ceil_div(page_per_tile, self.threads_per_warp) + + # atom_async_copy: async copy atom for page table load + atom_async_copy = cute.make_copy_atom( + cpasync.CopyG2SOp(cache_mode=cpasync.LoadCacheMode.ALWAYS), + cutlass.Int32, + num_bits_per_copy=cutlass.Int32.width, + ) + mPT_for_copy = cute.flat_divide(mPT, (1,)) + sPT_for_copy = cute.flat_divide(common_params.sPT, (1,)) + # elem_per_thread is a dynamic value depends on the page_size setting. + for i in range(elem_per_thread): + idx = i * self.threads_per_warp + tidx + if cute.elem_less( + k_index * page_per_tile + idx, mPT.shape[0] + ) and cute.elem_less(idx, page_per_tile): + cute.copy( + atom_async_copy, + mPT_for_copy[None, k_index * page_per_tile + idx], + sPT_for_copy[None, idx, load_pt_producer_state.index], + ) + else: + sPT_for_copy[None, idx, load_pt_producer_state.index].fill(0) + mbar_ptr = load_pt_pipeline.producer_get_barrier(load_pt_producer_state) + load_pt_pipeline.producer_commit(load_pt_producer_state) + load_pt_producer_state.advance() + k_index += 1 + k_tile_count -= 1 + + return load_pt_producer_state + + @cute.jit + def load_tma( + self, + common_params: SimpleNamespace, + qk_params: SimpleNamespace, + v_params: SimpleNamespace, + k_index: cutlass.Int32, + k_tile_count: cutlass.Int32, + load_q_producer_state: pipeline.PipelineState, + load_kv_producer_state: pipeline.PipelineState, + load_pt_consumer_state: pipeline.PipelineState, + load_pt_release_state: pipeline.PipelineState, + ) -> tuple[ + pipeline.PipelineState, + pipeline.PipelineState, + pipeline.PipelineState, + pipeline.PipelineState, + ]: + """Load wrap to load Q/C latent/rope tensors. Updates the load qkv producer state. + + :param common_params: The common parameters + :type common_params: SimpleNamespace + :param qk_params: The qk parameters + :type qk_params: SimpleNamespace + :param v_params: The v parameters + :type v_params: SimpleNamespace + :param k_index: The k index + :type k_index: cutlass.Int32 + :param k_tile_count: The k tile count + :type k_tile_count: cutlass.Int32 + :param load_q_producer_state: The load q producer state + :type load_q_producer_state: pipeline.PipelineState + :param load_kv_producer_state: The load kv producer state + :type load_kv_producer_state: pipeline.PipelineState + :param load_pt_consumer_state: The load pt consumer state + :type load_pt_consumer_state: pipeline.PipelineState + :param load_pt_release_state: The load pt release state + :type load_pt_release_state: pipeline.PipelineState + + :return: The load q producer state, load kv producer state, load pt consumer state, and load pt release state + :rtype: tuple[pipeline.PipelineState, pipeline.PipelineState, pipeline.PipelineState, pipeline.PipelineState] + """ + # page table + mPT = common_params.mPT[None, common_params.blk_coord[2]] + + # Flatten divide and partition global tensors for QK TMA load + # (bM, bK, rM, rK, rL) + mma_qk_tiler_mk = cute.select(self.mma_qk_tiler, mode=[0, 2]) + gQL = cute.flat_divide(qk_params.mQL, mma_qk_tiler_mk) + mma_qk_tiler_mk_rope = cute.select(self.mma_qk_rope_tiler, mode=[0, 2]) + gQR = cute.flat_divide(qk_params.mQR, mma_qk_tiler_mk_rope) + + thr_mma_qk = qk_params.tiled_mma_qk.get_slice( + common_params.blk_coord[0] % cute.size(qk_params.tiled_mma_qk.thr_id) + ) + tSgQL = thr_mma_qk.partition_A(gQL) + tSgQR = thr_mma_qk.partition_A(gQR) + + cta_m = min( + qk_params.tiled_mma_qk.op.shape_mnk[0] + // qk_params.tiled_mma_qk.thr_id.shape, + self.page_size, + ) + page_tile_size = min(self.page_size, cta_m) + gCL = cute.tiled_divide(qk_params.mCL, (page_tile_size, self.mma_qk_tiler[2])) + tSgCL = ( + gCL[ + None, + common_params.blk_coord[0] % qk_params.tiled_mma_qk.thr_id.shape, + None, + None, + ] + if cta_m < self.page_size + else gCL[None, 0, None, None] + ) + gKR = cute.tiled_divide(qk_params.mKR, (page_tile_size, self.mma_qk_tiler[2])) + tSgKR = ( + gKR[ + None, + common_params.blk_coord[0] % qk_params.tiled_mma_qk.thr_id.shape, + None, + None, + ] + if cta_m < self.page_size + else gKR[None, 0, None, None] + ) + + # tma partition for q, k latent/rope + # smem: ((atom_v, rest_v), STAGE) + # gmem: ((atom_v, rest_v), RestM, RestK, RestL) + tQsQ, tQLgQL_mkl = cpasync.tma_partition( + qk_params.tma_atom_q_latent, + 0, + cute.make_layout(1), + cute.group_modes(qk_params.sQ, 0, 3), + cute.group_modes(tSgQL, 0, 3), + ) + + tQsQ_rope, tQRgQR_mkl = cpasync.tma_partition( + qk_params.tma_atom_q_rope, + 0, + cute.make_layout(1), + cute.group_modes(qk_params.sQ_rope, 0, 3), + cute.group_modes(tSgQR, 0, 3), + ) + + tKCsKC, tCLgCL = cpasync.tma_partition( + qk_params.tma_atom_c_latent, + 0, + cute.make_layout(1), + qk_params.sKC, + tSgCL, + ) + + _, tKRgKR = cpasync.tma_partition( + qk_params.tma_atom_c_rope, + 0, + cute.make_layout(1), + qk_params.sKC, + tSgKR, + ) + + tQLgQL = tQLgQL_mkl[ + None, None, None, common_params.blk_coord[1], common_params.blk_coord[2] + ] + tQRgQR = tQRgQR_mkl[ + None, None, None, common_params.blk_coord[1], common_params.blk_coord[2] + ] + + # Flatten divide and partition global tensors for V TMA load + page_tile_size = min(self.page_size, self.mma_pv_tiler[2]) + gCLT = cute.flat_divide(v_params.mCLT, (self.mma_pv_tiler[1], page_tile_size)) + cta_n = self.mma_pv_tiler[1] // v_params.tiled_mma_pv.thr_id.shape + gCLT = cute.logical_divide(gCLT, (cta_n,))[ + (None, common_params.blk_coord[0]), None, None, None, None + ] + tOgCLT = cute.tiled_divide(gCLT, (cta_n, page_tile_size)) + tOgCLT = tOgCLT[None, 0, 0, None, None, None] + + # tma partition for vc + # smem: ((atom_v, rest_v), STAGE) + # gmem: ((atom_v, rest_v), RestM, RestK, RestL) + tVCsVC, tCLTgCLT = cpasync.tma_partition( + v_params.tma_atom_c_latent_transpose, + 0, + cute.make_layout(1), + v_params.sVC, + tOgCLT, + ) + + # set extra params + common_params.mPT = mPT + qk_params.tQLgQL = tQLgQL + qk_params.tQRgQR = tQRgQR + qk_params.tCLgCL = tCLgCL + qk_params.tKRgKR = tKRgKR + qk_params.tQsQ = tQsQ + qk_params.tQsQ_rope = tQsQ_rope + qk_params.tKCsKC = tKCsKC + v_params.tCLTgCLT = tCLTgCLT + v_params.tVCsVC = tVCsVC + + load_q_producer_state, load_kv_producer_state, load_pt_consumer_state = ( + self.load_tma_qk_one_k_tile( + common_params, + qk_params, + k_index, + k_tile_count, + load_q_producer_state, + load_kv_producer_state, + load_pt_consumer_state, + load_q=True, + ) + ) + k_index += 1 + k_tile_count -= 1 + while k_tile_count > 0: + load_q_producer_state, load_kv_producer_state, load_pt_consumer_state = ( + self.load_tma_qk_one_k_tile( + common_params, + qk_params, + k_index, + k_tile_count, + load_q_producer_state, + load_kv_producer_state, + load_pt_consumer_state, + load_q=False, + ) + ) + load_kv_producer_state, load_pt_release_state = self.load_tma_v_one_k_tile( + common_params, + v_params, + k_index - 1, + load_kv_producer_state, + load_pt_release_state, + ) + k_index += 1 + k_tile_count -= 1 + + # load last v tile + load_kv_producer_state, load_pt_release_state = self.load_tma_v_one_k_tile( + common_params, + v_params, + k_index - 1, + load_kv_producer_state, + load_pt_release_state, + ) + return ( + load_q_producer_state, + load_kv_producer_state, + load_pt_consumer_state, + load_pt_release_state, + ) + + @cute.jit + def load_tma_qk_one_k_tile( + self, + common_params: SimpleNamespace, + qk_params: SimpleNamespace, + k_index: cutlass.Int32, + k_tile_count: cutlass.Int32, + load_q_producer_state: pipeline.PipelineState, + load_kv_producer_state: pipeline.PipelineState, + load_pt_consumer_state: pipeline.PipelineState, + load_q: bool, + ) -> tuple[pipeline.PipelineState, pipeline.PipelineState, pipeline.PipelineState]: + """Load one k-tile of Q/C latent/rope tensors. Updates the load qkv producer state. + + :param common_params: The common parameters + :type common_params: SimpleNamespace + :param qk_params: The qk parameters + :type qk_params: SimpleNamespace + :param k_index: The k index + :type k_index: cutlass.Int32 + :param k_tile_count: The k tile count + :type k_tile_count: cutlass.Int32 + :param load_q_producer_state: The load q producer state + :type load_q_producer_state: pipeline.PipelineState + :param load_kv_producer_state: The load kv producer state + :type load_kv_producer_state: pipeline.PipelineState + :param load_pt_consumer_state: The load pt consumer state + :type load_pt_consumer_state: pipeline.PipelineState + :param load_q: Whether to load q + :type load_q: bool + + :return: The load q producer state, load kv producer state, and load pt consumer state + :rtype: tuple[pipeline.PipelineState, pipeline.PipelineState, pipeline.PipelineState] + """ + page_per_tile = ceil_div( + self.mma_qk_tiler[1] // self.page_size, qk_params.tiled_mma_qk.thr_id.shape + ) + common_params.load_pt_pipeline.consumer_wait(load_pt_consumer_state) + page_table_stage = load_pt_consumer_state.index + load_pt_consumer_state.advance() + k_idx = cute.make_rmem_tensor(cute.make_layout(page_per_tile), cutlass.Int32) + for i in cutlass.range_constexpr(page_per_tile): + k_idx[i] = ( + common_params.sPT[0, page_table_stage] + if self.mma_qk_tiler[1] // self.page_size == 1 + else common_params.sPT[ + i + common_params.blk_coord[0] * page_per_tile, page_table_stage + ] + ) + # load q once at first iteration + if cutlass.const_expr(load_q): + common_params.load_q_pipeline.producer_acquire(load_q_producer_state) + # get the mbar ptr from pipeline. + tma_bar_ptr = common_params.load_q_pipeline.producer_get_barrier( + load_q_producer_state + ) + for i in cutlass.range(self.iterations_qk_latent): + # load q latent + cute.copy( + qk_params.tma_atom_q_latent, + qk_params.tQLgQL[None, 0, i], + qk_params.tQsQ[None, (i, 0)], + tma_bar_ptr=tma_bar_ptr, + ) + for i in cutlass.range(self.iterations_qk_rope): + # load q rope + cute.copy( + qk_params.tma_atom_q_rope, + qk_params.tQRgQR[None, 0, i], + qk_params.tQsQ_rope[None, i], + tma_bar_ptr=tma_bar_ptr, + ) + load_q_producer_state.advance() + load_kv_pipeline = common_params.load_kv_pipeline + tma_bar_ptr = load_kv_pipeline.producer_get_barrier(load_kv_producer_state) + for i in cutlass.range(self.iterations_qk_latent): + # get the mbar ptr from pipeline. + tma_bar_ptr = load_kv_pipeline.producer_get_barrier(load_kv_producer_state) + load_kv_pipeline.producer_acquire(load_kv_producer_state) + for k in cutlass.range(page_per_tile): + # load k latent + cute.copy( + qk_params.tma_atom_c_latent, + qk_params.tCLgCL[None, i, k_idx[k]], + qk_params.tKCsKC[None, k, 0, load_kv_producer_state.index], + tma_bar_ptr=tma_bar_ptr, + ) + load_kv_producer_state.advance() + + for i in cutlass.range(self.iterations_qk_rope): + # get the mbar ptr from pipeline. + tma_bar_ptr = load_kv_pipeline.producer_get_barrier(load_kv_producer_state) + load_kv_pipeline.producer_acquire(load_kv_producer_state) + for k in cutlass.range(page_per_tile): + # load k rope + cute.copy( + qk_params.tma_atom_c_rope, + qk_params.tKRgKR[None, i, k_idx[k]], + qk_params.tKCsKC[None, k, 0, load_kv_producer_state.index], + tma_bar_ptr=tma_bar_ptr, + ) + load_kv_producer_state.advance() + + return load_q_producer_state, load_kv_producer_state, load_pt_consumer_state + + @cute.jit + def load_tma_v_one_k_tile( + self, + common_params: SimpleNamespace, + v_params: SimpleNamespace, + k_index: cutlass.Int32, + load_kv_producer_state: pipeline.PipelineState, + load_pt_release_state: pipeline.PipelineState, + ) -> tuple[pipeline.PipelineState, pipeline.PipelineState]: + """Load one k-tile of compressed latent transpose tensor(v). Updates the load qkv producer state. + + :param common_params: The common parameters + :type common_params: SimpleNamespace + :param v_params: The load tma v parameters + :type v_params: SimpleNamespace + :param k_index: The k index + :type k_index: cutlass.Int32 + :param load_kv_producer_state: The load qkv producer state + :type load_kv_producer_state: pipeline.PipelineState + :param load_pt_release_state: The load pt release state + :type load_pt_release_state: pipeline.PipelineState + + :return: The load kv producer state and load pt release state + :rtype: tuple[pipeline.PipelineState, pipeline.PipelineState] + """ + page_per_tile = self.mma_pv_tiler[2] * self.iterations_pv_k // self.page_size + page_per_subtile = ceil_div(page_per_tile, self.iterations_pv_k) + k_idx = cute.make_rmem_tensor(cute.make_layout(page_per_tile), cutlass.Int32) + page_table_stage = load_pt_release_state.index + for i in cutlass.range(page_per_tile): + k_idx[i] = ( + common_params.sPT[0, page_table_stage] + if page_per_tile == 1 + else common_params.sPT[i, page_table_stage] + ) + common_params.load_pt_pipeline.consumer_release(load_pt_release_state) + load_pt_release_state.advance() + load_kv_pipeline = common_params.load_kv_pipeline + tma_bar_ptr = load_kv_pipeline.producer_get_barrier(load_kv_producer_state) + for i in cutlass.range(self.iterations_pv_k): + for j in cutlass.range(self.iterations_pv_n): + # get the mbar ptr from pipeline. + tma_bar_ptr = load_kv_pipeline.producer_get_barrier( + load_kv_producer_state + ) + load_kv_pipeline.producer_acquire(load_kv_producer_state) + for k in cutlass.range(page_per_subtile): + k_idx_i = k_idx[ + k + + i + // ceil_div(self.iterations_pv_k, page_per_tile) + * page_per_subtile + ] + cute.copy( + v_params.tma_atom_c_latent_transpose, + v_params.tCLTgCLT[ + None, + j, + i % ceil_div(self.iterations_pv_k, page_per_tile), + k_idx_i, + ], + v_params.tVCsVC[None, 0, k, load_kv_producer_state.index], + tma_bar_ptr=tma_bar_ptr, + ) + + load_kv_producer_state.advance() + return load_kv_producer_state, load_pt_release_state + + @cute.jit + def mma( + self, + common_params: SimpleNamespace, + qk_params: SimpleNamespace, + pv_params: SimpleNamespace, + k_tile_count: cutlass.Int32, + tiled_mma_qk: cute.TiledMma, + tiled_mma_pv: cute.TiledMma, + load_q_consumer_state: pipeline.PipelineState, + load_kv_consumer_state: pipeline.PipelineState, + mma_s_producer_state: pipeline.PipelineState, + p_mma_consumer_state: pipeline.PipelineState, + mma_o_producer_state: pipeline.PipelineState, + ) -> tuple[ + cute.TiledMma, + cute.TiledMma, + pipeline.PipelineState, + pipeline.PipelineState, + pipeline.PipelineState, + pipeline.PipelineState, + ]: + """MMA warp to compute the result of Q*K^T and P*V. Updates the tiled mma and pipeline states. + + :param common_params: The common parameters for mma qk and pv + :type common_params: SimpleNamespace + :param qk_params: The mma qk parameters + :type qk_params: SimpleNamespace + :param pv_params: The mma pv parameters + :type pv_params: SimpleNamespace + :param k_tile_count: The k tile count + :type k_tile_count: cutlass.Int32 + :param tiled_mma_qk: The tiled mma qk + :type tiled_mma_qk: cute.TiledMma + :param tiled_mma_pv: The tiled mma pv + :type tiled_mma_pv: cute.TiledMma + :param load_q_consumer_state: The load q consumer state + :type load_q_consumer_state: pipeline.PipelineState + :param load_kv_consumer_state: The load kv consumer state + :type load_kv_consumer_state: pipeline.PipelineState + :param mma_s_producer_state: The mma s producer state + :type mma_s_producer_state: pipeline.PipelineState + :param p_mma_consumer_state: The p mma consumer state + :type p_mma_consumer_state: pipeline.PipelineState + :param mma_o_producer_state: The mma o producer state + :type mma_o_producer_state: pipeline.PipelineState + + :return: The tiled mma qk, the tiled mma pv, the load q consumer state, the load kv consumer state, the mma s producer state, the p mma consumer state, and the mma o producer state + :rtype: tuple[cute.TiledMma, cute.TiledMma, pipeline.PipelineState, pipeline.PipelineState, pipeline.PipelineState, pipeline.PipelineState, pipeline.PipelineState] + """ + + tSrQ = tiled_mma_qk.make_fragment_A(qk_params.sQ) + tSrQ_rope = tiled_mma_qk.make_fragment_A(qk_params.sQ_rope) + tSrKC = tiled_mma_qk.make_fragment_B(qk_params.sKC) + tOrP = tiled_mma_pv.make_fragment_A(pv_params.sP) + tOrVC = tiled_mma_pv.make_fragment_B(pv_params.sVC) + + tStS_shape = tiled_mma_qk.partition_shape_C( + cute.select(self.mma_qk_tiler, mode=[0, 1]) + ) + tStS_staged_fake = tiled_mma_qk.make_fragment_C( + cute.append(tStS_shape, self.mma_s_stage) + ) + # use real tmem ptr for tStS + tStS_staged = cute.make_tensor(common_params.tmem_ptr, tStS_staged_fake.layout) + tOtO_shape = tiled_mma_pv.partition_shape_C( + cute.select(self.mma_pv_tiler, mode=[0, 1]) + ) + # mma O has 1 stage. + tOtO = tiled_mma_pv.make_fragment_C(tOtO_shape) + tOtO_layout = cute.append( + tOtO.layout, + cute.make_layout( + common_params.L // self.mma_pv_tiler[1], + stride=self.mma_pv_tiler[1] // self.warps_in_n, + ), + ) + tOtO_staged = cute.make_tensor( + tStS_staged.iterator + self.tmem_o_offset, tOtO_layout + ) + + # set more parameters + qk_params.tSrQ = tSrQ + qk_params.tSrQ_rope = tSrQ_rope + qk_params.tSrKC = tSrKC + qk_params.tStS_staged = tStS_staged + pv_params.tOrP = tOrP + pv_params.tOrVC = tOrVC + pv_params.tOtO_staged = tOtO_staged + + # mma O accumulates on K, so the accumlate flag is set to False once before all K blocks. + tiled_mma_pv.set(tcgen05.Field.ACCUMULATE, False) + load_q_pipeline = common_params.load_q_pipeline + if common_params.is_leader_cta: + load_q_release_state = load_q_consumer_state.clone() + + ( + tiled_mma_qk, + load_q_consumer_state, + load_kv_consumer_state, + mma_s_producer_state, + ) = self.mma_qk( + common_params, + qk_params, + tiled_mma_qk, + load_q_consumer_state, + load_kv_consumer_state, + mma_s_producer_state, + wait_q=True, + ) + k_tile_count -= 1 + while k_tile_count > 0: + ( + tiled_mma_qk, + load_q_consumer_state, + load_kv_consumer_state, + mma_s_producer_state, + ) = self.mma_qk( + common_params, + qk_params, + tiled_mma_qk, + load_q_consumer_state, + load_kv_consumer_state, + mma_s_producer_state, + wait_q=False, + ) + ( + tiled_mma_pv, + load_kv_consumer_state, + p_mma_consumer_state, + mma_o_producer_state, + ) = self.mma_pv( + common_params, + pv_params, + tiled_mma_pv, + load_kv_consumer_state, + p_mma_consumer_state, + mma_o_producer_state, + ) + k_tile_count -= 1 + + # release q consumer states + load_q_pipeline.consumer_release(load_q_release_state) + load_q_release_state.advance() + ( + tiled_mma_pv, + load_kv_consumer_state, + p_mma_consumer_state, + mma_o_producer_state, + ) = self.mma_pv( + common_params, + pv_params, + tiled_mma_pv, + load_kv_consumer_state, + p_mma_consumer_state, + mma_o_producer_state, + ) + + return ( + tiled_mma_qk, + tiled_mma_pv, + load_q_consumer_state, + load_kv_consumer_state, + mma_s_producer_state, + p_mma_consumer_state, + mma_o_producer_state, + ) + + @cute.jit + def mma_qk( + self, + common_params: SimpleNamespace, + qk_params: SimpleNamespace, + tiled_mma_qk: cute.TiledMma, + load_q_consumer_state: pipeline.PipelineState, + load_kv_consumer_state: pipeline.PipelineState, + mma_s_producer_state: pipeline.PipelineState, + wait_q: bool, + ) -> tuple[ + cute.TiledMma, + pipeline.PipelineState, + pipeline.PipelineState, + pipeline.PipelineState, + ]: + """Compute one k-tile of mma for Q*K^T. Updates the tiled MMA QK and pipeline states. + + :param qk_params: The qk parameters + :type qk_params: SimpleNamespace + :param tiled_mma_qk: The tiled mma qk + :type tiled_mma_qk: cute.TiledMma + :param load_q_consumer_state: The load q consumer state + :type load_q_consumer_state: pipeline.PipelineState + :param load_kv_consumer_state: The load kv consumer state + :type load_kv_consumer_state: pipeline.PipelineState + :param mma_s_producer_state: The mma s producer state + :type mma_s_producer_state: pipeline.PipelineState + + :return: The tiled mma qk, the load q consumer state, the load kv consumer state, and the mma s producer state + :rtype: tuple[cute.TiledMma, pipeline.PipelineState, pipeline.PipelineState, pipeline.PipelineState] + """ + tStS = qk_params.tStS_staged[None, None, None, mma_s_producer_state.index] + + qk_params.mma_s_pipeline.producer_acquire(mma_s_producer_state) + tiled_mma_qk.set(tcgen05.Field.ACCUMULATE, False) + load_q_pipeline = common_params.load_q_pipeline + load_kv_pipeline = common_params.load_kv_pipeline + if cutlass.const_expr(wait_q): + load_q_pipeline.consumer_wait(load_q_consumer_state) + load_q_consumer_state.advance() + for q_stage in range(self.iterations_qk_latent): + load_kv_pipeline.consumer_wait(load_kv_consumer_state) + kc_stage = load_kv_consumer_state.index + for k_block in cutlass.range(cute.size(qk_params.tSrQ.shape[2])): + cute.gemm( + tiled_mma_qk, + tStS, + qk_params.tSrQ[None, None, k_block, q_stage], + qk_params.tSrKC[None, None, k_block, kc_stage], + tStS, + ) + tiled_mma_qk.set(tcgen05.Field.ACCUMULATE, True) + load_kv_pipeline.consumer_release(load_kv_consumer_state) + load_kv_consumer_state.advance() + for q_stage in range(self.iterations_qk_rope): + load_kv_pipeline.consumer_wait(load_kv_consumer_state) + kc_stage = load_kv_consumer_state.index + for k_block in cutlass.range(self.rope_dim // tiled_mma_qk.shape_mnk[2]): + cute.gemm( + tiled_mma_qk, + tStS, + qk_params.tSrQ_rope[None, None, k_block, q_stage], + qk_params.tSrKC[None, None, k_block, kc_stage], + tStS, + ) + tiled_mma_qk.set(tcgen05.Field.ACCUMULATE, True) + load_kv_pipeline.consumer_release(load_kv_consumer_state) + load_kv_consumer_state.advance() + + qk_params.mma_s_pipeline.producer_commit(mma_s_producer_state) + mma_s_producer_state.advance() + return ( + tiled_mma_qk, + load_q_consumer_state, + load_kv_consumer_state, + mma_s_producer_state, + ) + + @cute.jit + def mma_pv( + self, + common_params: SimpleNamespace, + pv_params: SimpleNamespace, + tiled_mma_pv: cute.TiledMma, + load_kv_consumer_state: pipeline.PipelineState, + p_mma_consumer_state: pipeline.PipelineState, + mma_o_producer_state: pipeline.PipelineState, + ) -> tuple[ + cute.TiledMma, + pipeline.PipelineState, + pipeline.PipelineState, + pipeline.PipelineState, + ]: + """Compute one k-tile of mma for P*V. Updates the tiled mma pv and pipeline states. + + :param common_params: The common parameters + :type common_params: SimpleNamespace + :param pv_params: The pv parameters + :type pv_params: SimpleNamespace + :param tiled_mma_pv: The tiled mma pv + :type tiled_mma_pv: cute.TiledMma + :param load_kv_consumer_state: The load kv consumer state + :type load_kv_consumer_state: pipeline.PipelineState + :param p_mma_consumer_state: The P MMA consumer state + :type p_mma_consumer_state: pipeline.PipelineState + :param mma_o_producer_state: The MMA o producer state + :type mma_o_producer_state: pipeline.PipelineState + + :return: The tiled mma pv, the load qkv consumer state, the P MMA consumer state, and the MMA o producer state + :rtype: tuple[cute.TiledMma, pipeline.PipelineState, pipeline.PipelineState, pipeline.PipelineState] + """ + + pv_params.mma_o_pipeline.producer_acquire(mma_o_producer_state) + pv_params.p_mma_pipeline.consumer_wait(p_mma_consumer_state) + load_kv_pipeline = common_params.load_kv_pipeline + for p_stage in range(self.iterations_pv_k): + accumulate_flag = tiled_mma_pv.get(tcgen05.Field.ACCUMULATE) + for acc_stage in range(self.iterations_pv_n): + load_kv_pipeline.consumer_wait(load_kv_consumer_state) + tiled_mma_pv.set(tcgen05.Field.ACCUMULATE, accumulate_flag) + vc_stage = load_kv_consumer_state.index + tOtO = pv_params.tOtO_staged[None, None, None, acc_stage] + for k_block in cutlass.range(pv_params.tOrP.shape[2]): + cute.gemm( + tiled_mma_pv, + tOtO, + pv_params.tOrP[ + None, + None, + k_block, + (p_stage, p_mma_consumer_state.index), + ], + pv_params.tOrVC[None, None, k_block, vc_stage], + tOtO, + ) + tiled_mma_pv.set(tcgen05.Field.ACCUMULATE, True) + load_kv_pipeline.consumer_release(load_kv_consumer_state) + load_kv_consumer_state.advance() + pv_params.p_mma_pipeline.consumer_release(p_mma_consumer_state) + p_mma_consumer_state.advance() + pv_params.mma_o_pipeline.producer_commit(mma_o_producer_state) + mma_o_producer_state.advance() + + return ( + tiled_mma_pv, + load_kv_consumer_state, + p_mma_consumer_state, + mma_o_producer_state, + ) + + @cute.jit + def compute( + self, + common_params: SimpleNamespace, + softmax_params: SimpleNamespace, + k_index: cutlass.Int32, + k_tile_count: cutlass.Int32, + mma_s_consumer_state: pipeline.PipelineState, + p_mma_producer_state: pipeline.PipelineState, + p_cor_producer_state: pipeline.PipelineState, + ) -> tuple[pipeline.PipelineState, pipeline.PipelineState, pipeline.PipelineState]: + """Compute warp to compute the result of softmax, rescale, and epilogue. Updates the related pipeline states. + + :param common_params: The common parameters + :type common_params: SimpleNamespace + :param softmax_params: The softmax parameters + :type softmax_params: SimpleNamespace + :param k_index: The index of the k-tile + :type k_index: cutlass.Int32 + :param k_tile_count: The number of k-tiles + :type k_tile_count: cutlass.Int32 + :param mma_s_consumer_state: The MMA s consumer state + :type mma_s_consumer_state: pipeline.PipelineState + :param p_mma_producer_state: The P MMA producer state + :type p_mma_producer_state: pipeline.PipelineState + :param p_cor_producer_state: The P correction producer state + :type p_cor_producer_state: pipeline.PipelineState + + :return: The MMA s consumer state, the P MMA producer state, and the P correction producer state + :rtype: tuple[pipeline.PipelineState, pipeline.PipelineState, pipeline.PipelineState] + """ + + k_tile_total = cute.ceil_div(common_params.K, self.mma_qk_tiler[1]) + + row_max = -self.acc_dtype.inf + row_sum = self.acc_dtype(0) + correction_factor = self.acc_dtype(1) + common_params.p_cor_pipeline.producer_acquire(p_cor_producer_state) + + # no mask applied + while k_tile_count > 1: + ( + mma_s_consumer_state, + p_mma_producer_state, + p_cor_producer_state, + row_max, + row_sum, + correction_factor, + ) = self.softmax( + common_params, + softmax_params, + k_index, + mma_s_consumer_state, + p_mma_producer_state, + p_cor_producer_state, + row_max, + row_sum, + correction_factor, + False, + False, + ) + k_index = k_index + 1 + k_tile_count = k_tile_count - 1 + + # mask applied + if cutlass.const_expr(common_params.mAccO is not None): + ( + mma_s_consumer_state, + p_mma_producer_state, + p_cor_producer_state, + row_max, + row_sum, + correction_factor, + ) = self.softmax( + common_params, + softmax_params, + k_index, + mma_s_consumer_state, + p_mma_producer_state, + p_cor_producer_state, + row_max, + row_sum, + correction_factor, + k_index == k_tile_total - 1, + True, + ) + else: + ( + mma_s_consumer_state, + p_mma_producer_state, + p_cor_producer_state, + row_max, + row_sum, + correction_factor, + ) = self.softmax( + common_params, + softmax_params, + k_index, + mma_s_consumer_state, + p_mma_producer_state, + p_cor_producer_state, + row_max, + row_sum, + correction_factor, + True, + True, + ) + + return mma_s_consumer_state, p_mma_producer_state, p_cor_producer_state + + @cute.jit + def correction( + self, + common_params: SimpleNamespace, + epilogue_params: SimpleNamespace, + k_tile_count: cutlass.Int32, + p_cor_consumer_state: pipeline.PipelineState, + mma_o_consumer_state: pipeline.PipelineState, + ) -> tuple[pipeline.PipelineState, pipeline.PipelineState]: + """Compute warp to compute the result of softmax, rescale, and epilogue. Updates the related pipeline states. + + :param common_params: The common parameters + :type common_params: SimpleNamespace + :param epilogue_params: The epilogue parameters + :type epilogue_params: SimpleNamespace + :param k_index: The index of the k-tile + :type k_index: cutlass.Int32 + :param k_tile_count: The number of k-tiles + :type k_tile_count: cutlass.Int32 + :param p_cor_consumer_state: The P correction consumer state + :type p_cor_consumer_state: pipeline.PipelineState + :param mma_o_consumer_state: The MMA o consumer state + :type mma_o_consumer_state: pipeline.PipelineState + + :return: The P correction consumer state, and the MMA o consumer state + :rtype: tuple[pipeline.PipelineState, pipeline.PipelineState] + """ + + k_tile_count_init = k_tile_count + while k_tile_count > 0: + p_cor_consumer_state, row_sum, row_max, correction_factor, no_correction = ( + self.get_correction_factor(common_params, p_cor_consumer_state) + ) + if k_tile_count_init != k_tile_count: + mma_o_consumer_state = self.rescale( + common_params, + mma_o_consumer_state, + correction_factor, + no_correction, + ) + k_tile_count = k_tile_count - 1 + if k_tile_count == 0: + mma_o_consumer_state = self.epilogue( + common_params, + epilogue_params, + mma_o_consumer_state, + row_sum, + row_max, + ) + + return p_cor_consumer_state, mma_o_consumer_state + + @cute.jit + def exchange_p_cor_metadata( + self, + common_params: SimpleNamespace, + softmax_params: SimpleNamespace, + correction_factor: cutlass.Float32, + row_sum: cutlass.Float32, + row_max: cutlass.Float32, + row_max_new: cutlass.Float32, + tAcc: cute.Tensor, + tidx: cutlass.Int32, + p_cor_producer_state: pipeline.PipelineState, + ) -> pipeline.PipelineState: + """Compute the correction factor for the last k tile.""" + no_correction = 0 + if ( + row_max_new - row_max + ) * softmax_params.softmax_scale_log2 <= self.skip_correction_threshold: + no_correction = 1 + row_max_new = row_max + + # pad for 4x32b + corr_layout = cute.make_layout( + (tAcc.shape[0], (4, tAcc.shape[1][1]), self.mma_s_stage), + stride=(tAcc.stride[0], (1, tAcc.stride[1][1]), 4), + ) + tCor = cute.make_tensor( + common_params.tmem_ptr + self.correction_factor_offset, + corr_layout, + ) + cCor = cute.make_identity_tensor(tCor.shape) + corr_tmem_store_atom = cute.make_copy_atom( + tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(4)), self.acc_dtype + ) + corr_tmem_store_tiled_copy = tcgen05.make_tmem_copy(corr_tmem_store_atom, tCor) + corr_tmem_store_thr_copy = corr_tmem_store_tiled_copy.get_slice(tidx) + cCor_for_copy = corr_tmem_store_thr_copy.partition_S(cCor) + tCor_for_copy = corr_tmem_store_thr_copy.partition_D(tCor) + rCor = cute.make_fragment_like( + cCor_for_copy[None, None, None, 0], self.acc_dtype + ) + rCor_int = cute.make_tensor( + cute.recast_ptr(rCor.iterator, dtype=cutlass.Int32), rCor.layout + ) + rCor[0] = row_sum + rCor[1] = row_max_new + rCor[2] = correction_factor + rCor_int[3] = no_correction + + cute.copy( + corr_tmem_store_tiled_copy, + rCor, + tCor_for_copy[None, None, None, p_cor_producer_state.index], + ) + # fence between tmem store and correction warp + cute.arch.fence_view_async_tmem_store() + common_params.p_cor_pipeline.producer_commit(p_cor_producer_state) + p_cor_producer_state.advance() + return p_cor_producer_state, row_max_new + + @cute.jit + def softmax( + self, + common_params: SimpleNamespace, + softmax_params: SimpleNamespace, + k_index: cutlass.Int32, + mma_s_consumer_state: pipeline.PipelineState, + p_mma_producer_state: pipeline.PipelineState, + p_cor_producer_state: pipeline.PipelineState, + row_max: cutlass.Float32, + row_sum: cutlass.Float32, + correction_factor: cutlass.Float32, + is_last_tile: bool, + is_local_last_tile: cutlass.Boolean, + ) -> tuple[ + pipeline.PipelineState, + pipeline.PipelineState, + pipeline.PipelineState, + cutlass.Float32, + cutlass.Float32, + cutlass.Float32, + ]: + """Softmax for one k-tile. Updates the related pipeline states and returns the computed results. + + :param common_params: The common parameters + :type common_params: SimpleNamespace + :param softmax_params: The softmax parameters + :type softmax_params: SimpleNamespace + :param k_index: The index of the k-tile + :type k_index: cutlass.Int32 + :param mma_s_consumer_state: The MMA s consumer state + :type mma_s_consumer_state: pipeline.PipelineState + :param p_mma_producer_state: The P MMA producer state + :type p_mma_producer_state: pipeline.PipelineState + :param p_cor_producer_state: The P correction producer state + :type p_cor_producer_state: pipeline.PipelineState + :param row_max: The row max + :type row_max: cutlass.Float32 + :param row_sum: The row sum + :type row_sum: cutlass.Float32 + :param correction_factor: The correction factor + :type correction_factor: cutlass.Float32 + :param is_last_tile: Whether the last tile + :type is_last_tile: bool + :param is_local_last_tile: Whether the last tile is local + :type is_local_last_tile: cutlass.Boolean + + :return: The MMA s consumer state, the P MMA producer state, the P correction producer state, the row max, the row sum, and the correction factor + :rtype: tuple[pipeline.PipelineState, pipeline.PipelineState, pipeline.PipelineState, cutlass.Float32, cutlass.Float32, cutlass.Float32] + """ + + softmax_params.p_mma_pipeline.producer_acquire(p_mma_producer_state) + softmax_params.mma_s_pipeline.consumer_wait(mma_s_consumer_state) + + # load S from tmem + tStS_shape = softmax_params.tiled_mma_qk.partition_shape_C( + cute.select(self.mma_qk_tiler, mode=[0, 1]) + ) + tStS_staged_fake = softmax_params.tiled_mma_qk.make_fragment_C( + cute.append(tStS_shape, self.mma_s_stage) + ) + tStS_staged = cute.make_tensor(common_params.tmem_ptr, tStS_staged_fake.layout) + tStS = tStS_staged[None, None, None, mma_s_consumer_state.index] + + tAcc = tStS[(None, None), 0, 0] + cta_qk_tiler = ( + self.mma_qk_tiler[0] // self.cluster_shape_mnk[0], + self.mma_qk_tiler[1], + self.mma_qk_tiler[2], + ) + cS = cute.make_identity_tensor(cute.select(cta_qk_tiler, mode=[0, 1])) + + tmem_load_atom = cute.make_copy_atom( + tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(32)), self.acc_dtype + ) + tmem_tiled_copy = tcgen05.make_tmem_copy(tmem_load_atom, tAcc) + + tidx = common_params.tidx % (self.num_compute_warps * self.threads_per_warp) + + tmem_thr_copy = tmem_tiled_copy.get_slice(tidx) + tTR_tAcc = tmem_thr_copy.partition_S(tAcc) + tTR_tS = tmem_thr_copy.partition_D(cS) + + tTR_rAcc = cute.make_fragment_like(tTR_tS, self.acc_dtype) + + row_max_new = row_max + arch = BaseDSL._get_dsl().get_arch_enum() + if cutlass.const_expr(arch >= Arch.sm_100 and arch <= Arch.sm_100f): + cute.copy(tmem_tiled_copy, tTR_tAcc, tTR_rAcc) + for i in cutlass.range_constexpr(cute.size(tTR_rAcc)): + if is_last_tile: + tTR_rAcc[i] = ( + tTR_rAcc[i] + if cute.elem_less( + tTR_tS[i][1] + self.mma_qk_tiler[1] * k_index, + common_params.K, + ) + else -self.acc_dtype.inf + ) + # reduction for row_max + row_max_new = tTR_rAcc.load().reduce(cute.ReductionOp.MAX, row_max_new, 0) + + elif cutlass.const_expr(arch >= Arch.sm_103 and arch <= Arch.sm_103f): + tmem_load_red_atom = cute.make_copy_atom( + tcgen05.copy.LdRed32x32bOp( + tcgen05.copy.Repetition(64), redOp=tcgen05.TmemLoadRedOp.MAX + ), + self.acc_dtype, + ) + tmem_red_tiled_copy = tcgen05.make_tmem_copy(tmem_load_red_atom, tAcc) + tmem_red_thr_copy = tmem_red_tiled_copy.get_slice(tidx) + tTR_tAcc_red = tmem_red_thr_copy.partition_S(tAcc) + tTR_tS_red = tmem_red_thr_copy.partition_D(cS) + tTR_rAcc_red = cute.make_fragment_like(tTR_tS_red, self.acc_dtype) + tTR_rMax = cute.make_rmem_tensor( + cute.make_layout((1, tTR_tS_red.shape[1], tTR_tS_red.shape[2])), + self.acc_dtype, + ) + cute.copy( + tmem_red_tiled_copy, + tTR_tAcc_red, + (tTR_rAcc_red, tTR_rMax), + ) + tTR_rAcc = cute.make_tensor(tTR_rAcc_red.iterator, tTR_rAcc.layout) + if is_last_tile: + for i in cutlass.range_constexpr(cute.size(tTR_rAcc)): + tTR_rAcc[i] = ( + tTR_rAcc[i] + if cute.elem_less( + tTR_tS[i][1] + self.mma_qk_tiler[1] * k_index, + common_params.K, + ) + else -self.acc_dtype.inf + ) + # reduction for row_max + row_max_new = tTR_rAcc.load().reduce( + cute.ReductionOp.MAX, row_max_new, 0 + ) + else: + row_max_new = cute.arch.fmax(row_max_new, tTR_rMax[0]) + + # if warps in N is 2, reduce row_max across warps (0, 1) and (2, 3) + if cutlass.const_expr(self.warps_in_n == 2): + common_params.smem_exchange[tidx] = row_max_new + self.softmax_exchange_sync_bar.wait() + row_max_new = cute.arch.fmax( + row_max_new, + common_params.smem_exchange[ + (tidx + 64) % (self.num_compute_warps * self.threads_per_warp) + ], + ) + + # find correction factor + correction_factor = cute.math.exp2( + (row_max - row_max_new) * softmax_params.softmax_scale_log2, fastmath=True + ) + # split kv case + if cutlass.const_expr(not is_local_last_tile): + p_cor_producer_state, row_max_new = self.exchange_p_cor_metadata( + common_params, + softmax_params, + correction_factor, + row_sum, + row_max, + row_max_new, + tAcc, + tidx, + p_cor_producer_state, + ) + + # softmax + fma_b = softmax_params.softmax_scale_log2 + fma_c = (0.0 - row_max_new) * softmax_params.softmax_scale_log2 + + for i in cutlass.range(cute.size(tTR_rAcc), vectorize=True, unroll_full=True): + tTR_rAcc[i] = tTR_rAcc[i] * fma_b + fma_c + tTR_rAcc[i] = cute.math.exp2(tTR_rAcc[i], fastmath=True) + + tTR_rS = cute.make_fragment_like(tTR_tS, self.q_dtype) + + # quantize + tTR_rS.store(tTR_rAcc.load().to(self.q_dtype)) + + # create sP + sP = softmax_params.sP[None, None, None, (None, p_mma_producer_state.index)] + sP_mk_view = cute.make_tensor( + sP.iterator, + cute.make_layout( + ( + (sP.shape[0][0], sP.shape[1]), + (sP.shape[0][1], sP.shape[2], sP.shape[3]), + ), + stride=( + (sP.stride[0][0], sP.stride[1]), + (sP.stride[0][1], sP.stride[2], sP.stride[3]), + ), + ), + ) + # change to PISL + sP_wo_swizzle_iter = cute.recast_ptr(sP.iterator, swizzle_=None) + swizzle_bits = ( + int(math.log2(self.mma_pv_tiler[2] * self.q_dtype.width // 8 // 32)) + 1 + ) + swizzle_base = 3 if self.q_dtype.width == 16 else 4 + sP_swizzle = cute.make_swizzle(swizzle_bits, swizzle_base, 3) + sP_mk_view = cute.make_tensor( + sP_wo_swizzle_iter, + cute.make_composed_layout(sP_swizzle, 0, sP_mk_view.layout), + ) + universal_copy_bits = 128 + smem_copy_atom = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), + self.q_dtype, + num_bits_per_copy=universal_copy_bits, + ) + smem_tiled_copy = cute.make_tiled_copy_D(smem_copy_atom, tmem_tiled_copy) + smem_thr_copy = smem_tiled_copy.get_slice(tidx) + rP_copy_view = smem_thr_copy.retile(tTR_rS) + sP_copy_view = smem_thr_copy.partition_D(sP_mk_view) + cute.copy(smem_tiled_copy, rP_copy_view, sP_copy_view) + + # fence between smem store and mma o + cute.arch.fence_view_async_shared() + softmax_params.p_mma_pipeline.producer_commit(p_mma_producer_state) + p_mma_producer_state.advance() + + # row_sum, using `add_packed_f32x2` to reduce the number of instructions + row_sum = row_sum * correction_factor + row_sum_vec = (0.0, 0.0) + for i in cutlass.range_constexpr(0, cute.size(tTR_rAcc), 2): + row_sum_vec = cute.arch.add_packed_f32x2( + row_sum_vec, (tTR_rAcc[i], tTR_rAcc[i + 1]) + ) + row_sum = row_sum_vec[0] + row_sum_vec[1] + row_sum + + # split kv case + if cutlass.const_expr(is_local_last_tile): + p_cor_producer_state, row_max_new = self.exchange_p_cor_metadata( + common_params, + softmax_params, + correction_factor, + row_sum, + row_max, + row_max_new, + tAcc, + tidx, + p_cor_producer_state, + ) + + # store correction factor/row_sum/row_max to tmem for correction warp + common_params.p_cor_pipeline.producer_acquire(p_cor_producer_state) + + # fence between tmem load and mma s + cute.arch.fence_view_async_tmem_load() + + softmax_params.mma_s_pipeline.consumer_release(mma_s_consumer_state) + mma_s_consumer_state.advance() + + return ( + mma_s_consumer_state, + p_mma_producer_state, + p_cor_producer_state, + row_max_new, + row_sum, + correction_factor, + ) + + @cute.jit + def _tmem_load_partition( + self, common_params: SimpleNamespace, tiled_mma_pv: cute.TiledMma, iter_n: int + ) -> tuple[ + cute.TiledMma, cute.TiledMma, cute.TiledMma, cute.TiledMma, cute.TiledMma + ]: + """Tensor memory load partition for rescale and epilogue. + + :param common_params: The common parameters + :type common_params: SimpleNamespace + :param tiled_mma_pv: The tiled mma pv + :type tiled_mma_pv: cute.TiledMma + :param iter_n: The iteration number + :type iter_n: int + + :return: The tiled mma pv, the tiled mma pv, the tiled mma pv, the tiled mma pv, the tiled mma pv + :rtype: tuple[cute.TiledMma, cute.TiledMma, cute.TiledMma, cute.TiledMma, cute.TiledMma] + """ + + tOtO_shape = tiled_mma_pv.partition_shape_C( + cute.select(self.mma_pv_tiler, mode=[0, 1]) + ) + tOtO = tiled_mma_pv.make_fragment_C(tOtO_shape) + tOtO_layout = cute.append( + tOtO.layout, + cute.make_layout( + common_params.L // self.mma_pv_tiler[1], + stride=self.mma_pv_tiler[1] // self.warps_in_n, + ), + ) + tOtO = cute.make_tensor( + common_params.tmem_ptr + self.tmem_o_offset, tOtO_layout + ) + tOtO = tOtO[None, None, None, iter_n] + + tAcc = tOtO[(None, None), 0, 0] + + tmem_load_atom = cute.make_copy_atom( + tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(32)), self.acc_dtype + ) + tmem_load_tiled_copy = tcgen05.make_tmem_copy(tmem_load_atom, tAcc) + tmem_load_thr_copy = tmem_load_tiled_copy.get_slice( + common_params.tidx % (self.num_compute_warps * self.threads_per_warp) + ) + + cta_pv_tiler = ( + self.mma_pv_tiler[0] // self.cluster_shape_mnk[0], + self.mma_pv_tiler[1], + self.mma_pv_tiler[2], + ) + # Flatten divide and partition global tensors for O + cta_pv_tiler_mn = cute.select(cta_pv_tiler, mode=[0, 1]) + + gO = None + if cutlass.const_expr(common_params.mAccO is not None): + gO = cute.local_tile( + common_params.mAccO[None, common_params.blk_coord[3], None, None, None], + cta_pv_tiler_mn, + ( + common_params.blk_coord[0], + iter_n, + common_params.blk_coord[1], + common_params.blk_coord[2], + ), + ) + cO = cute.local_tile( + cute.make_identity_tensor( + common_params.mAccO[ + None, common_params.blk_coord[3], None, None, None + ].shape + ), + cta_pv_tiler_mn, + ( + common_params.blk_coord[0], + iter_n, + common_params.blk_coord[1], + common_params.blk_coord[2], + ), + ) + else: + gO = cute.local_tile( + common_params.mO, + cta_pv_tiler_mn, + ( + common_params.blk_coord[0], + iter_n, + common_params.blk_coord[1], + common_params.blk_coord[2], + ), + ) + cO = cute.local_tile( + cute.make_identity_tensor(common_params.mO.shape), + cta_pv_tiler_mn, + ( + common_params.blk_coord[0], + iter_n, + common_params.blk_coord[1], + common_params.blk_coord[2], + ), + ) + tTR_tAcc = tmem_load_thr_copy.partition_S(tAcc) + tTR_gO = tmem_load_thr_copy.partition_D(gO) + tTR_cO = tmem_load_thr_copy.partition_D(cO) + tTR_rAcc = cute.make_fragment_like(tTR_gO, self.acc_dtype) + return tmem_load_tiled_copy, tAcc, tTR_tAcc, tTR_gO, tTR_cO, tTR_rAcc + + def get_correction_factor( + self, + common_params: SimpleNamespace, + p_cor_consumer_state: pipeline.PipelineState, + ) -> tuple[ + pipeline.PipelineState, + cutlass.Float32, + cutlass.Float32, + cutlass.Float32, + cutlass.Int32, + ]: + """Get the correction factor from the P correction consumer state. + + :param common_params: The common parameters + :type common_params: SimpleNamespace + :param p_cor_consumer_state: The P correction consumer state + :type p_cor_consumer_state: pipeline.PipelineState + + :return: The P correction consumer state, the row_sum, the row_max, and the correction factor + :rtype: tuple[pipeline.PipelineState, cutlass.Float32, cutlass.Float32, cutlass.Float32, cutlass.Int32] + """ + common_params.p_cor_pipeline.consumer_wait(p_cor_consumer_state) + tidx = common_params.tidx % (self.num_compute_warps * self.threads_per_warp) + # load correction factor + _, tAcc, _, _, _, _ = self._tmem_load_partition( + common_params, common_params.tiled_mma_pv, 0 + ) + corr_layout = cute.make_layout( + (tAcc.shape[0], (4, tAcc.shape[1][1]), self.p_cor_stage), + stride=(tAcc.stride[0], (1, tAcc.stride[1][1]), 4), + ) + tCor = cute.make_tensor( + common_params.tmem_ptr + self.correction_factor_offset, corr_layout + ) + cCor = cute.make_identity_tensor(tCor.shape) + corr_tmem_load_atom = cute.make_copy_atom( + tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(4)), self.acc_dtype + ) + corr_tmem_load_tiled_copy = tcgen05.make_tmem_copy(corr_tmem_load_atom, tCor) + corr_tmem_load_thr_copy = corr_tmem_load_tiled_copy.get_slice(tidx) + tCor_for_copy = corr_tmem_load_thr_copy.partition_S(tCor) + cCor_for_copy = corr_tmem_load_thr_copy.partition_D(cCor) + rCor = cute.make_fragment_like( + cCor_for_copy[None, None, None, 0], self.acc_dtype + ) + rCor_int = cute.make_tensor( + cute.recast_ptr(rCor.iterator, dtype=cutlass.Int32), rCor.layout + ) + cute.copy( + corr_tmem_load_tiled_copy, + tCor_for_copy[None, None, None, p_cor_consumer_state.index], + rCor, + ) + row_sum = rCor[0] + row_max = rCor[1] + correction_factor = rCor[2] + no_correction = rCor_int[3] + + common_params.p_cor_pipeline.consumer_release(p_cor_consumer_state) + p_cor_consumer_state.advance() + return p_cor_consumer_state, row_sum, row_max, correction_factor, no_correction + + @cute.jit + def rescale( + self, + common_params: SimpleNamespace, + mma_o_consumer_state: pipeline.PipelineState, + correction_factor: cutlass.Float32, + no_correction: cutlass.Int32, + ) -> pipeline.PipelineState: + """Rescale for one k-tile. Updates the related pipeline state. + + :param common_params: The common parameters + :type common_params: SimpleNamespace + :param mma_o_consumer_state: The mma o consumer state + :type mma_o_consumer_state: pipeline.PipelineState + :param correction_factor: The correction factor + :type correction_factor: cutlass.Float32 + :param no_correction: Whether to apply correction factor + :type no_correction: cutlass.Int32 + + :return: The MMA o consumer state + :rtype: pipeline.PipelineState + """ + skip_correction = cute.arch.vote_all_sync(no_correction == 1) + common_params.mma_o_pipeline.consumer_wait(mma_o_consumer_state) + if not skip_correction: + for iter_n in cutlass.range_constexpr(self.iterations_pv_n): + # tmem load tiled copy and partition results. + tmem_load_tiled_copy, tAcc, tTR_tAcc, tTR_gO, tTR_cO, tTR_rAcc = ( + self._tmem_load_partition( + common_params, common_params.tiled_mma_pv, iter_n + ) + ) + + # tmem store tiled copy + tmem_store_atom = cute.make_copy_atom( + tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(32)), self.acc_dtype + ) + tmem_store_tiled_copy = tcgen05.make_tmem_copy(tmem_store_atom, tAcc) + + # load o + cute.copy(tmem_load_tiled_copy, tTR_tAcc, tTR_rAcc) + # rescale, using `mul_packed_f32x2` to reduce the number of instructions + for i in cutlass.range( + cute.size(tTR_rAcc), vectorize=True, unroll_full=True + ): + tTR_rAcc[i] = tTR_rAcc[i] * correction_factor + + # store o to tensor memory for next k tile + cute.copy(tmem_store_tiled_copy, tTR_rAcc, tTR_tAcc) + + cute.arch.fence_view_async_tmem_store() + common_params.mma_o_pipeline.consumer_release(mma_o_consumer_state) + mma_o_consumer_state.advance() + + return mma_o_consumer_state + + @cute.jit + def epilogue( + self, + common_params: SimpleNamespace, + epilogue_params: SimpleNamespace, + mma_o_consumer_state: pipeline.PipelineState, + row_sum: cutlass.Float32, + row_max: cutlass.Float32, + ) -> pipeline.PipelineState: + """Epilogue for one k-tile. Updates the related pipeline state. + + :param common_params: The common parameters + :type common_params: SimpleNamespace + :param epilogue_params: The epilogue parameters + :type epilogue_params: SimpleNamespace + :param mma_o_consumer_state: The mma o consumer state + :type mma_o_consumer_state: pipeline.PipelineState + :param row_sum: The row sum + :type row_sum: cutlass.Float32 + :param row_max: The row max + :type row_max: cutlass.Float32 + + :return: The MMA o consumer state + :rtype: pipeline.PipelineState + """ + + tidx = common_params.tidx % (self.num_compute_warps * self.threads_per_warp) + + # exchange row_sum between warps (0, 1) and (2, 3) + if cutlass.const_expr(self.warps_in_n == 2): + common_params.smem_exchange[tidx] = row_sum + self.epilogue_exchange_sync_bar.wait() + # (64, 2) + row_sum = ( + row_sum + + common_params.smem_exchange[ + (tidx + 64) % (self.num_compute_warps * self.threads_per_warp) + ] + ) + # mma_o pipeline consumer wait + common_params.mma_o_pipeline.consumer_wait(mma_o_consumer_state) + for iter_n in cutlass.range_constexpr(self.iterations_pv_n): + # tmem load tiled copy and partition results. + tmem_load_tiled_copy, tAcc, tTR_tAcc, tTR_gO, tTR_cO, tTR_rAcc = ( + self._tmem_load_partition( + common_params, common_params.tiled_mma_pv, iter_n + ) + ) + + # load o + cute.copy(tmem_load_tiled_copy, tTR_tAcc, tTR_rAcc) + + # apply output scale and normalize by row_sum + for i in cutlass.range( + cute.size(tTR_rAcc), vectorize=True, unroll_full=True + ): + tTR_rAcc[i] = ( + tTR_rAcc[i] + * epilogue_params.output_scale + * cute.arch.rcp_approx(row_sum) + ) + + # store o to global memory + tR2G_rO_src = None + tR2G_rO_dst = tTR_gO + if cutlass.const_expr(common_params.mAccO is None): + tR2G_rO_src = cute.make_fragment_like(tTR_gO, self.o_dtype) + # using final output dtype for o + tR2G_rO_src.store(tTR_rAcc.load().to(self.o_dtype)) + else: + # using accumulate dtype for o + tR2G_rO_src = tTR_rAcc + + if cute.elem_less(tTR_cO[0][0], common_params.H): + cute.autovec_copy( + tR2G_rO_src, + tR2G_rO_dst, + l1c_evict_priority=cute.nvgpu.CacheEvictionPriority.NO_ALLOCATE, + ) + + # store the lse to global memory + cta_pv_tiler = ( + self.mma_pv_tiler[0] // self.cluster_shape_mnk[0], + self.mma_pv_tiler[1], + self.mma_pv_tiler[2], + ) + gLSE = None + cLSE = None + if cutlass.const_expr(epilogue_params.mAccLSE is None): + gLSE = cute.local_tile( + epilogue_params.mLSE, + (cta_pv_tiler[0], 1, 1), + ( + common_params.blk_coord[0], + common_params.blk_coord[1], + common_params.blk_coord[2], + ), + (1, 1, 1), + ) + cLSE = cute.local_tile( + cute.make_identity_tensor(epilogue_params.mLSE.shape), + (cta_pv_tiler[0], 1, 1), + ( + common_params.blk_coord[0], + common_params.blk_coord[1], + common_params.blk_coord[2], + ), + (1, 1, 1), + ) + + else: + gLSE = cute.local_tile( + epilogue_params.mAccLSE[ + None, common_params.blk_coord[3], None, None + ], + (cta_pv_tiler[0], 1, 1), + ( + common_params.blk_coord[0], + common_params.blk_coord[1], + common_params.blk_coord[2], + ), + (1, 1, 1), + ) + cLSE = cute.local_tile( + cute.make_identity_tensor( + epilogue_params.mAccLSE[ + None, common_params.blk_coord[3], None, None + ].shape + ), + (cta_pv_tiler[0], 1, 1), + ( + common_params.blk_coord[0], + common_params.blk_coord[1], + common_params.blk_coord[2], + ), + (1, 1, 1), + ) + lse = ( + cute.math.log2(row_sum, fastmath=True) + + epilogue_params.softmax_scale_log2 * row_max + ) + if cutlass.const_expr(self.warps_in_n == 2): + if cute.elem_less(cLSE[tidx][0], common_params.H): + gLSE[tidx] = lse + + cute.arch.fence_view_async_tmem_load() + common_params.mma_o_pipeline.consumer_release(mma_o_consumer_state) + mma_o_consumer_state.advance() + + return mma_o_consumer_state + + def make_and_init_load_pt_pipeline(self, load_pt_mbar_ptr): + """Create and initialize the load page table pipeline. + + :param load_pt_mbar_ptr: The load page table mbar pointer + :type load_pt_mbar_ptr: cute.Tensor + + :return: The load page table pipeline + :rtype: pipeline.PipelineAsync + """ + load_pt_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + self.threads_per_warp * len([self.load_pt_warp_id]), + ) + load_pt_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + self.threads_per_warp * len([self.load_tma_warp_id]), + ) + return pipeline.PipelineCpAsync.create( + barrier_storage=load_pt_mbar_ptr, + num_stages=self.load_pt_stage, + producer_group=load_pt_producer_group, + consumer_group=load_pt_consumer_group, + defer_sync=True, + ) + + def make_and_init_load_qkv_pipeline( + self, load_qkv_mbar_ptr, cta_layout_vmnk, load_stages, tx_count + ) -> pipeline.PipelineTmaUmma: + """Create and initialize the tma load qkv pipeline. + + :param load_qkv_mbar_ptr: The load qkv mbar pointer + :type load_qkv_mbar_ptr: cute.Tensor + :param cta_layout_vmnk: The cta layout vmnk + :type cta_layout_vmnk: tuple[int, int, int] + :param load_stages: The load stages + :type load_stages: list[int] + :param tx_count: The tx count + :type tx_count: int + + :return: The tma load qkv pipeline + :rtype: pipeline.PipelineTmaUmma + """ + load_qkv_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, len([self.load_tma_warp_id]) + ) + load_qkv_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, len([self.mma_warp_id]) + ) + return pipeline.PipelineTmaUmma.create( + barrier_storage=load_qkv_mbar_ptr, + num_stages=load_stages, + producer_group=load_qkv_producer_group, + consumer_group=load_qkv_consumer_group, + tx_count=tx_count, + cta_layout_vmnk=cta_layout_vmnk, + defer_sync=True, + ) + + def make_and_init_mma_s_pipeline( + self, mma_s_mbar_ptr, cta_layout_vmnk + ) -> pipeline.PipelineUmmaAsync: + """Create and initialize the mma s pipeline. + + :param mma_s_mbar_ptr: The mma s mbar pointer + :type mma_s_mbar_ptr: cute.Tensor + :param cta_layout_vmnk: The cta layout vmnk + :type cta_layout_vmnk: tuple[int, int, int] + + :return: The mma s pipeline + :rtype: pipeline.PipelineUmmaAsync + """ + + mma_s_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, len([self.mma_warp_id]) + ) + consumer_thread_size = ( + self.threads_per_warp + * len(self.compute_warp_ids) + * self.cluster_shape_mnk[0] + ) + mma_s_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + consumer_thread_size, + ) + return pipeline.PipelineUmmaAsync.create( + barrier_storage=mma_s_mbar_ptr, + num_stages=self.mma_s_stage, + producer_group=mma_s_producer_group, + consumer_group=mma_s_consumer_group, + cta_layout_vmnk=cta_layout_vmnk, + defer_sync=True, + ) + + def make_and_init_p_mma_pipeline( + self, p_mma_mbar_ptr, cta_layout_vmnk + ) -> pipeline.PipelineAsyncUmma: + """Create and initialize the p mma pipeline. + + :param p_mma_mbar_ptr: The p mma mbar pointer + :type p_mma_mbar_ptr: cute.Tensor + :param cta_layout_vmnk: The cta layout vmnk + :type cta_layout_vmnk: tuple[int, int, int] + + :return: The p mma pipeline + :rtype: pipeline.PipelineAsyncUmma + """ + + producer_thread_size = ( + self.threads_per_warp + * len(self.compute_warp_ids) + * self.cluster_shape_mnk[0] + ) + p_mma_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + producer_thread_size, + ) + p_mma_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, len([self.mma_warp_id]) + ) + return pipeline.PipelineAsyncUmma.create( + barrier_storage=p_mma_mbar_ptr, + num_stages=self.p_mma_stage, + producer_group=p_mma_producer_group, + consumer_group=p_mma_consumer_group, + cta_layout_vmnk=cta_layout_vmnk, + defer_sync=True, + ) + + def make_and_init_p_cor_pipeline( + self, p_cor_mbar_ptr + ) -> pipeline.PipelineAsyncUmma: + """Create and initialize the p correction pipeline. + + :param p_cor_mbar_ptr: The p correction mbar pointer + :type p_cor_mbar_ptr: cute.Tensor + + :return: The p correction pipeline + :rtype: pipeline.PipelineAsyncUmma + """ + + producer_thread_size = self.threads_per_warp * len(self.compute_warp_ids) + p_cor_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + producer_thread_size, + ) + p_cor_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + producer_thread_size, + ) + return pipeline.PipelineAsync.create( + barrier_storage=p_cor_mbar_ptr, + num_stages=self.p_cor_stage, + producer_group=p_cor_producer_group, + consumer_group=p_cor_consumer_group, + defer_sync=True, + ) + + def make_and_init_mma_o_pipeline( + self, mma_o_mbar_ptr, cta_layout_vmnk + ) -> pipeline.PipelineUmmaAsync: + """Create and initialize the mma o pipeline. + + :param mma_o_mbar_ptr: The mma o mbar pointer + :type mma_o_mbar_ptr: cute.Tensor + :param cta_layout_vmnk: The cta layout vmnk + :type cta_layout_vmnk: tuple[int, int, int] + + :return: The mma o pipeline + :rtype: pipeline.PipelineUmmaAsync + """ + + mma_o_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, len([self.mma_warp_id]) + ) + consumer_thread_size = ( + self.threads_per_warp + * len(self.compute_warp_ids) + * self.cluster_shape_mnk[0] + ) + mma_o_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + consumer_thread_size, + ) + return pipeline.PipelineUmmaAsync.create( + barrier_storage=mma_o_mbar_ptr, + num_stages=self.mma_o_stage, + producer_group=mma_o_producer_group, + consumer_group=mma_o_consumer_group, + cta_layout_vmnk=cta_layout_vmnk, + defer_sync=True, + ) + + @staticmethod + def _compute_grid( + o: cute.Tensor, + split_kv: cutlass.Int32, + cluster_shape_mnk: Tuple[int, int, int], + max_active_clusters: int, + is_persistent: bool, + ) -> Tuple[MLAStaticTileSchedulerParams, Tuple[int, int, int]]: + """Compute grid shape for the output tensor C. + + :param c: The output tensor C + :type c: cute.Tensor + :param cta_tile_shape_mnk: The shape (M, N, K) of the CTA tile. + :type cta_tile_shape_mnk: tuple[int, int, int] + :param cluster_shape_mn: Shape of each cluster in M, N dimensions. + :type cluster_shape_mn: tuple[int, int] + + :return: Tile scheduler parameters and grid shape. + :rtype: tuple[MLAStaticTileSchedulerParams, tuple[int, int, int]] + """ + o_shape = o.shape + tile_sched_params = create_mla_static_tile_scheduler_params( + is_persistent, + cute.size(o_shape[3]), + cute.size(o_shape[2]), + cluster_shape_mnk, + split_kv, + ) + grid = MLAStaticTileScheduler.get_grid_shape( + tile_sched_params, max_active_clusters + ) + + return tile_sched_params, grid + + @staticmethod + def get_workspace_size( + H: int, + S: int, + D: int, + B: int, + split_kv: int, + acc_dtype: Type[cutlass.Numeric], + ) -> int: + """Get the extra workspace(device memory) size for the MLA kernel when split_kv is not 1. + + :param H: The height of the output tensor C + :type H: int + :param S: The sequence length of the output tensor C + :type S: int + :param D: The depth of the output tensor C + :type D: int + :param B: The batch size of the output tensor C + :type B: int + :param split_kv: The split key-value of the output tensor C + :type split_kv: int + :param acc_dtype: The data type of the output tensor C + :type acc_dtype: Type[cutlass.Numeric] + + :return: The workspace size for the MLA kernel + :rtype: int + """ + if split_kv == 1: + return 0 + return B * H * S * split_kv * (D + 1) * acc_dtype.width // 8 + + @cute.jit + def initialize_workspace( + self, + H: cutlass.Int32, + D: cutlass.Int32, + S: cutlass.Int32, + B: cutlass.Int32, + split_kv: cutlass.Int32, + acc_dtype: Type[cutlass.Numeric], + workspace: cute.Tensor, + ) -> tuple[cute.Tensor, cute.Tensor]: + """Initialize the workspace for the MLA kernel. Construct the intermediate tensors + acc_o and acc_lse. + + :param H: The height of the output tensor C + :type H: cutlass.Int32 + :param D: The depth of the output tensor C + :type D: cutlass.Int32 + :param S: The sequence length of the output tensor C + :type S: cutlass.Int32 + :param B: The batch size of the output tensor C + :type B: cutlass.Int32 + :param split_kv: The split key-value of the output tensor C + :type split_kv: cutlass.Int32 + :param acc_dtype: The data type of the output tensor C + :type acc_dtype: Type[cutlass.Numeric] + :param workspace: The workspace tensor + :type workspace: cute.Tensor + + :return: The output tensor C and the workspace tensor + :rtype: tuple[cute.Tensor, cute.Tensor] + """ + acc_o, acc_lse = None, None + if cutlass.const_expr(workspace is not None): + align = 256 // self.q_dtype.width + acc_o_layout = cute.make_layout( + (H, split_kv, D, S, B), + stride=( + cute.assume(split_kv * D, align), + cute.assume(D, align), + 1, + cute.assume(split_kv * H * D, align), + cute.assume(H * split_kv * S * D, align), + ), + ) + acc_o_iter = cute.recast_ptr(workspace.iterator, dtype=acc_dtype) + acc_o = cute.make_tensor(acc_o_iter, acc_o_layout) + acc_lse_layout = cute.make_layout( + (H, split_kv, S, B), + stride=(split_kv, 1, H * split_kv, H * split_kv * S), + ) + acc_lse_iter = cute.recast_ptr( + workspace.iterator + cute.cosize(acc_o_layout) * acc_dtype.width // 8, + dtype=acc_dtype, + ) + acc_lse = cute.make_tensor(acc_lse_iter, acc_lse_layout) + return acc_o, acc_lse + + @staticmethod + def can_implement( + B: int, + S: int, + K: int, + H: int, + L: int, + R: int, + in_dtype: Type[cutlass.Numeric], + out_dtype: Type[cutlass.Numeric], + acc_dtype: Type[cutlass.Numeric], + lse_dtype: Type[cutlass.Numeric], + mma_qk_tiler_mn: Tuple[int, int], + mma_pv_tiler_mn: Tuple[int, int], + split_kv: int, + is_persistent: bool, + is_var_seq: bool, + is_var_split_kv: bool, + page_size: int, + ) -> bool: + """Check if the MLA kernel can be implemented. + + :param B: The batch size of the output tensor C + :type B: int + :param S: The sequence length of the output tensor C + :type S: int + :param K: The width of the output tensor KV + :type K: int + :param H: The number of heads of the output tensor C + :type H: int + :param L: The number of latent dimensions of the tensor KV + :type L: int + :param R: The number of rope dimensions of the tensor C_rope + :type R: int + :param in_dtype: The data type of the input tensor + :type in_dtype: Type[cutlass.Numeric] + :param out_dtype: The data type of the output tensor + :type out_dtype: Type[cutlass.Numeric] + :param acc_dtype: The data type of the accumulator + :type acc_dtype: Type[cutlass.Numeric] + :param lse_dtype: The data type of the log-sum-exp + :type lse_dtype: Type[cutlass.Numeric] + :param mma_qk_tiler_mn: The tile shape of the query-key matrix multiplication + :type mma_qk_tiler_mn: Tuple[int, int] + :param mma_pv_tiler_mn: The tile shape of the probability-value matrix multiplication + :type mma_pv_tiler_mn: Tuple[int, int] + :param split_kv: The split key-value of the output tensor C + :type split_kv: int + :param is_persistent: Whether to use persistent kernel optimization + :type is_persistent: bool + :param is_var_seq: Whether to use variable sequence length + :type is_var_seq: bool + :param is_var_split_kv: Whether to use variable split_kv + :type is_var_split_kv: bool + :param page_size: The page size of the page table + :type page_size: int + + :return: Whether the MLA kernel can be implemented + :rtype: bool + """ + if L != 512 or R != 64: + return False + if in_dtype not in [cutlass.Float16]: + return False + if out_dtype not in [cutlass.Float16]: + return False + if acc_dtype != cutlass.Float32 or lse_dtype != cutlass.Float32: + return False + # page size equals 1 is prohibited by tma specification, not 128B aligned. + if mma_qk_tiler_mn[1] % page_size != 0 or page_size == 1: + return False + if mma_qk_tiler_mn[0] != mma_pv_tiler_mn[0] or mma_qk_tiler_mn[0] != 128: + return False + if is_var_split_kv and not is_var_seq: + return False + if H > 128 or (H < 128 and split_kv != 1): + return False + if S < 1 or S > 4: + return False + if K <= 0: + return False + return True + + +def run( + batch_size: int, + seq_len_q: int, + seq_len_k: int, + num_heads: int, + latent_dim: int, + rope_dim: int, + in_dtype: Type[cutlass.Numeric], + out_dtype: Type[cutlass.Numeric], + acc_dtype: Type[cutlass.Numeric], + lse_dtype: Type[cutlass.Numeric], + mma_qk_tiler_mn: Tuple[int, int], + mma_pv_tiler_mn: Tuple[int, int], + split_kv: int, + is_persistent: bool, + is_var_seq: bool, + is_var_split_kv: bool, + page_size: int, + softmax_scale: float, + output_scale: float, + skip_correction_threshold: float, + tolerance: float, + warmup_iterations: int, + iterations: int, + skip_ref_check: bool, + use_cold_l2: bool, + **kwargs, +): + """Execute Multi-Head Latent Attention (MLA) on Blackwell architecture and validate results. + + This function creates random input tensors for query latent/rope, compressed latent/rope, and value, + then performs the complete MLA computation pipeline. It supports configurable data types, tiling parameters, + page table, variable sequence length, and variable split_kv. Results can be validated against a PyTorch reference + implementation or run multiple times for performance measurement. + + :param batch_size: Batch size + :type batch_size: int + :param seq_len_q: Sequence length of Q + :type seq_len_q: int + :param seq_len_k: Sequence length of K + :type seq_len_k: int + :param num_heads: Number of heads + :type num_heads: int + :param latent_dim: dimension of query/compressed latent + :type latent_dim: int + :param rope_dim: dimension of query/compressed rope + :type rope_dim: int + :param in_dtype: Input data type for query/compressed latent/rope tensors + :type in_dtype: Type[cutlass.Numeric] + :param out_dtype: Output data type for attention output + :type out_dtype: Type[cutlass.Numeric] + :param acc_dtype: Accumulator data type for query-key matrix multiplication + :type acc_dtype: Type[cutlass.Numeric] + :param lse_dtype: Accumulator data type for log-sum-exp + :type lse_dtype: Type[cutlass.Numeric] + :param mma_qk_tiler_mn: Matrix multiply accumulate tile shape (M, N) for query-key matrix multiplication + :type mma_qk_tiler_mn: Tuple[int, int] + :param mma_pv_tiler_mn: Matrix multiply accumulate tile shape (M, N) for probability-value matrix multiplication + :type mma_pv_tiler_mn: Tuple[int, int] + :param split_kv: Split key-value + :type split_kv: int + :param is_persistent: Whether to use persistent kernel optimization + :type is_persistent: bool + :param is_var_seq: Whether to use variable sequence length + :type is_var_seq: bool + :param is_var_split_kv: Whether to use variable split_kv + :type is_var_split_kv: bool + :param page_size: Page size of the page table + :type page_size: int + :param softmax_scale: Attention score scaling factor + :type softmax_scale: float + :param output_scale: Output scaling factor + :type output_scale: float + :param skip_correction_threshold: Threshold to skip correction + :type skip_correction_threshold: float + :param tolerance: Maximum acceptable error for validation + :type tolerance: float + :param warmup_iterations: Number of warmup iterations + :type warmup_iterations: int + :param iterations: Number of iterations to run for performance testing + :type iterations: int + :param skip_ref_check: Skip validation against reference implementation + :type skip_ref_check: bool + :param use_cold_l2: Whether to use cold L2 cache + :type use_cold_l2: bool + + :raises ValueError: If input shapes are incompatible or head dimension is unsupported + :raises RuntimeError: If GPU is unavailable for computation + """ + + print("Running Blackwell MLA test with:") + print(f" batch_size: {batch_size}") + print(f" seq_len_q: {seq_len_q}") + print(f" seq_len_k: {seq_len_k}") + print(f" num_heads: {num_heads}") + print(f" latent_dim: {latent_dim}") + print(f" rope_dim: {rope_dim}") + print(f" in_dtype: {in_dtype}") + print(f" out_dtype: {out_dtype}") + print(f" acc_dtype: {acc_dtype}") + print(f" mma_qk_tiler_mn: {mma_qk_tiler_mn}") + print(f" mma_pv_tiler_mn: {mma_pv_tiler_mn}") + print(f" split_kv: {split_kv}") + print(f" is_persistent: {is_persistent}") + print(f" is_var_seq: {is_var_seq}") + print(f" is_var_split_kv: {is_var_split_kv}") + print(f" page_size: {page_size}") + print(f" softmax_scale: {softmax_scale}") + print(f" output_scale: {output_scale}") + print(f" skip_correction_threshold: {skip_correction_threshold}") + print(f" tolerance: {tolerance}") + print(f" warmup_iterations: {warmup_iterations}") + print(f" iterations: {iterations}") + print(f" skip_ref_check: {skip_ref_check}") + print(f" use_cold_l2: {use_cold_l2}") + + import torch + import cutlass.torch as cutlass_torch + + # Prepare pytorch tensors: Q, K, V (random from 0 to 2) and O (all zero) + if not torch.cuda.is_available(): + raise RuntimeError("GPU is required to run this example!") + + if not BlackwellMultiHeadLatentAttentionForwardFP16.can_implement( + batch_size, + seq_len_q, + seq_len_k, + num_heads, + latent_dim, + rope_dim, + in_dtype, + out_dtype, + acc_dtype, + lse_dtype, + mma_qk_tiler_mn, + mma_pv_tiler_mn, + split_kv, + is_persistent, + is_var_seq, + is_var_split_kv, + page_size, + ): + raise TypeError( + f"Unsupported testcase {batch_size}, {seq_len_q}, {seq_len_k}, {num_heads}, {latent_dim}, {rope_dim}, {in_dtype}, {out_dtype}, {acc_dtype}, {lse_dtype}, {mma_qk_tiler_mn}, {mma_pv_tiler_mn}, {split_kv}, {is_persistent}, {is_var_seq}, {is_var_split_kv}, {page_size}" + ) + + torch.manual_seed(1111) + + def create_data_tensor( + B, + HK, + D, + dtype, + is_dynamic_layout=True, + page_table=None, + cache_seqs=None, + is_lse=False, + seq_len_q=None, + ): + shape = (B, HK, D) + if page_table is not None: + if cache_seqs is not None: + max_seq_len = torch.max(cache_seqs) + shape = (B * ceil_div(max_seq_len, page_size), page_size, D) + else: + shape = (B * ceil_div(HK, page_size), page_size, D) + + if seq_len_q is not None: + shape = (B, seq_len_q, HK, D) + + permute_order = (1, 2, 0) + stride_order = (2, 0, 1) + leading_dim = 1 + if is_lse: + shape = (B, seq_len_q, HK) + permute_order = (2, 1, 0) + stride_order = (2, 1, 0) + leading_dim = 0 + elif seq_len_q is not None: + permute_order = (2, 3, 1, 0) + stride_order = (3, 2, 0, 1) + leading_dim = 1 + + init_config = cutlass.torch.RandomInitConfig(min_val=-2, max_val=2) + + torch_dtype = ( + cutlass_torch.dtype(dtype) if dtype != cutlass.Float8E4M3FN else torch.int8 + ) + + # Create dtype torch tensor (cpu) + torch_tensor_cpu = cutlass_torch.create_and_permute_torch_tensor( + shape, + torch_dtype, + permute_order=permute_order, + init_type=cutlass.torch.TensorInitType.RANDOM, + init_config=init_config, + ) + + # Create dtype torch tensor (gpu) + torch_tensor_gpu = torch_tensor_cpu.cuda() + + # Create f32 torch tensor (cpu) + f32_torch_tensor = torch_tensor_cpu.to(dtype=torch.float32) + + # Create dtype cute tensor (gpu) + cute_tensor = from_dlpack(torch_tensor_gpu, assumed_align=16) + cute_tensor.element_type = dtype + if is_dynamic_layout: + cute_tensor = cute_tensor.mark_layout_dynamic(leading_dim=leading_dim) + if not is_lse: + cute_tensor = cute_tensor.mark_compact_shape_dynamic( + mode=leading_dim, + stride_order=stride_order, + divisibility=(128 // dtype.width), + ) + + cute_tensor = cutlass_torch.convert_cute_tensor( + f32_torch_tensor, + cute_tensor, + dtype, + is_dynamic_layout=is_dynamic_layout, + ) + + return f32_torch_tensor, cute_tensor, torch_tensor_gpu + + def create_cache_seqs(batch_size, seq_len_k, is_var_seq): + cache_seqs_ref = torch.ones(batch_size, dtype=torch.int32) * seq_len_k + cache_seqs_gpu = cache_seqs_ref.cuda() + cache_seqs = from_dlpack(cache_seqs_gpu, assumed_align=16).mark_layout_dynamic() + if is_var_seq: + max_seq_len = seq_len_k + min_seq_len = int(seq_len_k * 0.8) + cache_seqs_ref = cutlass_torch.create_and_permute_torch_tensor( + (batch_size,), + torch.int32, + init_type=cutlass.torch.TensorInitType.RANDOM, + init_config=cutlass.torch.RandomInitConfig( + min_val=min_seq_len, max_val=max_seq_len + 1 + ), + ) + cache_seqs_gpu = cache_seqs_ref.cuda() + cache_seqs = from_dlpack( + cache_seqs_gpu, + assumed_align=16, + ).mark_layout_dynamic() + return cache_seqs_ref, cache_seqs, cache_seqs_gpu + + def create_page_table(batch_size, seq_len_k, is_var_seq, page_size): + max_seq_len = seq_len_k if not is_var_seq else torch.max(cache_seqs_ref) + page_count = ceil_div(max_seq_len, page_size) + page_table_ref = torch.empty([batch_size, page_count], dtype=torch.int32) + # use transposed index for page table to make sure the value is in bound of `batch_size * seq_len_block`. In practice, the value could be any positive values. This setting is only for testing purpose. + for b in range(batch_size): + for j in range(page_count): + page_table_ref[b, j] = b + j * batch_size + page_table_gpu = page_table_ref.permute(1, 0).cuda() + page_table = from_dlpack(page_table_gpu, assumed_align=16).mark_layout_dynamic( + leading_dim=0 + ) + return page_table_ref, page_table, page_table_gpu + + def create_block_split_kvs( + batch_size, + split_kv, + cache_seqs_ref, + is_var_split_kv, + mma_qk_tiler_mn, + cluster_shape_mnk, + max_active_clusters, + ): + block_split_kvs_ref, block_split_kvs, block_split_kvs_gpu = None, None, None + # check if split_kv is valid otherwise do auto setting of split_kv + if is_var_split_kv: + block_split_kvs_ref = torch.zeros([batch_size], dtype=torch.int32) + for b in range(batch_size): + block_split_kvs_ref[b] = ( + BlackwellMultiHeadLatentAttentionForwardFP16.get_split_kv( + batch_size, + seq_len_q, + cache_seqs_ref[b].item(), + mma_qk_tiler_mn, + max_active_clusters * cluster_shape_mnk[0], + ) + ) + split_kv = torch.max(block_split_kvs_ref).item() + block_split_kvs_gpu = block_split_kvs_ref.cuda() + block_split_kvs = from_dlpack( + block_split_kvs_gpu, assumed_align=16 + ).mark_layout_dynamic() + elif split_kv <= 0: + split_kv = BlackwellMultiHeadLatentAttentionForwardFP16.get_split_kv( + batch_size, + seq_len_q, + cache_seqs_ref[0].item(), + mma_qk_tiler_mn, + max_active_clusters * cluster_shape_mnk[0], + ) + return split_kv, block_split_kvs_ref, block_split_kvs, block_split_kvs_gpu + + def create_workspace( + num_heads, seq_len_q, latent_dim, batch_size, split_kv, acc_dtype + ): + workspace_size = ( + BlackwellMultiHeadLatentAttentionForwardFP16.get_workspace_size( + num_heads, + seq_len_q, + latent_dim, + batch_size, + split_kv, + acc_dtype, + ) + ) + + workspace, workspace_torch = None, None + if workspace_size > 0: + workspace_torch = torch.empty([workspace_size], dtype=torch.int8).cuda() + workspace = from_dlpack(workspace_torch, assumed_align=32) + return workspace, workspace_torch + + cache_seqs_ref, cache_seqs, cache_seqs_torch = create_cache_seqs( + batch_size, seq_len_k, is_var_seq + ) + page_table_ref, page_table, page_table_torch = create_page_table( + batch_size, seq_len_k, is_var_seq, page_size + ) + cluster_shape_mnk = (2, 1, 1) + hardware_info = utils.HardwareInfo() + max_active_clusters = hardware_info.get_max_active_clusters( + cluster_shape_mnk[0] * cluster_shape_mnk[1] + ) + split_kv, block_split_kvs_ref, block_split_kvs, block_split_kvs_torch = ( + create_block_split_kvs( + batch_size, + split_kv, + cache_seqs_ref, + is_var_split_kv, + mma_qk_tiler_mn, + cluster_shape_mnk, + max_active_clusters, + ) + ) + + q_latent_ref, q_latent, q_latent_torch = create_data_tensor( + batch_size, + num_heads, + latent_dim, + in_dtype, + is_dynamic_layout=True, + seq_len_q=seq_len_q, + ) + q_rope_ref, q_rope, q_rope_torch = create_data_tensor( + batch_size, + num_heads, + rope_dim, + in_dtype, + is_dynamic_layout=True, + seq_len_q=seq_len_q, + ) + + c_latent_ref, c_latent, c_latent_torch = create_data_tensor( + batch_size, + seq_len_k, + latent_dim, + in_dtype, + is_dynamic_layout=True, + page_table=page_table, + cache_seqs=cache_seqs_ref, + ) + c_rope_ref, c_rope, c_rope_torch = create_data_tensor( + batch_size, + seq_len_k, + rope_dim, + in_dtype, + is_dynamic_layout=True, + page_table=page_table, + cache_seqs=cache_seqs_ref, + ) + o_ref, o, o_torch = create_data_tensor( + batch_size, + num_heads, + latent_dim, + out_dtype, + is_dynamic_layout=True, + seq_len_q=seq_len_q, + ) + lse_ref, lse, lse_torch = create_data_tensor( + batch_size, + num_heads, + 1, + lse_dtype, + is_dynamic_layout=True, + is_lse=True, + seq_len_q=seq_len_q, + ) + workspace, workspace_torch = create_workspace( + num_heads, seq_len_q, latent_dim, batch_size, split_kv, acc_dtype + ) + + mla = BlackwellMultiHeadLatentAttentionForwardFP16( + acc_dtype, + lse_dtype, + mma_qk_tiler_mn, + mma_pv_tiler_mn, + max_active_clusters, + page_size, + skip_correction_threshold, + is_persistent, + is_var_seq, + is_var_split_kv, + ) + + # Get current CUDA stream from PyTorch + torch_stream = torch.cuda.current_stream() + # Get the raw stream pointer as a CUstream + stream = cuda.CUstream(torch_stream.cuda_stream) + + # compile mla kernel + compiled_mla = cute.compile( + mla, + q_latent, + q_rope, + c_latent, + c_rope, + page_table, + o, + lse, + workspace, + split_kv, + cache_seqs, + block_split_kvs, + softmax_scale, + output_scale, + stream, + options="--opt-level 2", + ) + + def torch_reference_mla( + q_latent, + q_rope, + c_latent, + c_rope, + page_table, + cache_seqs, + softmax_scale=1.0, + output_scale=1.0, + ): + # expand and concat q_latent and q_rope to have the dimension of sequence length for q + q_ref = torch.cat([q_latent, q_rope], dim=1).permute(3, 2, 0, 1) + # expand and concat c_latent and c_rope to have the dimension of num_heads for k and v + page_count = page_table_ref.shape[1] + k_ref_paged = ( + torch.cat([c_latent, c_rope], dim=1) + .permute(2, 0, 1) + .reshape(batch_size * page_count, page_size, latent_dim + rope_dim) + ) + v_ref_paged = c_latent.permute(2, 0, 1).reshape( + batch_size * page_count, page_size, latent_dim + ) + + if is_var_seq: + max_seq_len = torch.max(cache_seqs_ref) + else: + max_seq_len = seq_len_k + + k_ref = torch.zeros([batch_size, 1, max_seq_len, latent_dim + rope_dim]) + v_ref = torch.zeros([batch_size, 1, max_seq_len, latent_dim]) + k_ref = torch.index_select( + k_ref_paged, 0, torch.flatten(page_table_ref) + ).reshape(batch_size, 1, -1, latent_dim + rope_dim)[:, :, :max_seq_len, :] + v_ref = torch.index_select( + v_ref_paged, 0, torch.flatten(page_table_ref) + ).reshape(batch_size, 1, -1, latent_dim)[:, :, :max_seq_len, :] + for b in range(batch_size): + k_ref[b, :, cache_seqs_ref[b] :, :] = 0 + v_ref[b, :, cache_seqs_ref[b] :, :] = 0 + import torch.nn.functional as F + + o_ref = F.scaled_dot_product_attention( + q_ref, + k_ref, + v_ref, + attn_mask=None, + dropout_p=0.0, + scale=softmax_scale, + is_causal=False, + ) + s_ref = torch.einsum("bhld,bhsd->bhls", q_ref, k_ref) + s_ref_max, s_ref_max_pos = torch.max(s_ref, dim=-1, keepdim=True) + softmax_scale_log2 = LOG2_E * softmax_scale + s_ref_sum = torch.sum( + torch.exp2((s_ref - s_ref_max) * softmax_scale_log2), dim=-1, keepdim=True + ) + + lse_ref = s_ref_max * softmax_scale_log2 + torch.log2(s_ref_sum) + lse_ref = lse_ref.squeeze(3).permute(2, 1, 0) + o_ref = o_ref * output_scale + o_ref = o_ref.permute(2, 3, 1, 0) + + return o_ref, lse_ref + + if skip_correction_threshold > 0.0: + print( + "Skipping correction verification since skip_correction_threshold is greater than 0.0..." + ) + skip_ref_check = True + if not skip_ref_check: + # Execute kernel once for reference checking + compiled_mla( + q_latent, + q_rope, + c_latent, + c_rope, + page_table, + o, + lse, + workspace, + split_kv, + cache_seqs, + block_split_kvs, + softmax_scale, + output_scale, + stream, + ) + torch.cuda.synchronize() + + print("Verifying results...") + if in_dtype == cutlass.Float8E4M3FN: + tolerance = 0.13 + o_ref, lse_ref = torch_reference_mla( + q_latent_ref, + q_rope_ref, + c_latent_ref, + c_rope_ref, + page_table, + cache_seqs, + softmax_scale, + output_scale, + ) + + if out_dtype in [cutlass.Float8E5M2, cutlass.Float8E4M3FN]: + # convert o back to f32 for comparison + o_fp32, o_fp32_torch = cutlass_torch.cute_tensor_like( + torch.empty(*o_torch.shape, dtype=torch.float32), + cutlass.Float32, + is_dynamic_layout=True, + assumed_align=16, + ) + cute.testing.convert(o, o_fp32) + o = o_fp32_torch.cpu() + ref_fp8, _ = cutlass_torch.cute_tensor_like( + torch.empty( + *o_ref.permute(3, 2, 0, 1).shape, dtype=torch.uint8 + ).permute(2, 3, 1, 0), + out_dtype, + is_dynamic_layout=True, + assumed_align=16, + ) + o_ref_gpu = o_ref.cuda() + o_ref_f32 = from_dlpack(o_ref_gpu).mark_layout_dynamic(leading_dim=1) + + # convert ref : f32 -> fp8 -> f32 + cute.testing.convert(o_ref_f32, ref_fp8) + cute.testing.convert(ref_fp8, o_ref_f32) + + o_ref = o_ref_gpu.cpu() + else: + o = o_torch.cpu().to(torch.float32) + lse = lse_torch.cpu() + lse_ref = lse_ref.to(cutlass.torch.dtype(lse_dtype)) + # Assert close results + torch.testing.assert_close(o, o_ref, atol=tolerance, rtol=1e-05) + torch.testing.assert_close(lse, lse_ref, atol=tolerance, rtol=1e-05) + print("Results verified successfully!") + + def generate_tensors(): + _, cache_seqs, _ = create_cache_seqs(batch_size, seq_len_k, is_var_seq) + _, page_table, _ = create_page_table( + batch_size, seq_len_k, is_var_seq, page_size + ) + _split_kv, _, block_split_kvs, _ = create_block_split_kvs( + batch_size, + split_kv, + cache_seqs_ref, + is_var_split_kv, + mma_qk_tiler_mn, + cluster_shape_mnk, + max_active_clusters, + ) + + _, q_latent, _ = create_data_tensor( + batch_size, + num_heads, + latent_dim, + in_dtype, + is_dynamic_layout=True, + seq_len_q=seq_len_q, + ) + _, q_rope, _ = create_data_tensor( + batch_size, + num_heads, + rope_dim, + in_dtype, + is_dynamic_layout=True, + seq_len_q=seq_len_q, + ) + + _, c_latent, _ = create_data_tensor( + batch_size, + seq_len_k, + latent_dim, + in_dtype, + is_dynamic_layout=True, + page_table=page_table, + cache_seqs=cache_seqs_ref, + ) + _, c_rope, _ = create_data_tensor( + batch_size, + seq_len_k, + rope_dim, + in_dtype, + is_dynamic_layout=True, + page_table=page_table, + cache_seqs=cache_seqs_ref, + ) + _, o, _ = create_data_tensor( + batch_size, + num_heads, + latent_dim, + out_dtype, + is_dynamic_layout=True, + seq_len_q=seq_len_q, + ) + _, lse, _ = create_data_tensor( + batch_size, + num_heads, + 1, + lse_dtype, + is_dynamic_layout=True, + is_lse=True, + seq_len_q=seq_len_q, + ) + workspace, workspace_torch = create_workspace( + num_heads, seq_len_q, latent_dim, batch_size, _split_kv, acc_dtype + ) + return testing.JitArguments( + q_latent, + q_rope, + c_latent, + c_rope, + page_table, + o, + lse, + workspace, + _split_kv, + cache_seqs, + block_split_kvs, + softmax_scale, + output_scale, + stream, + ) + + workspace_count = 1 + if use_cold_l2: + one_workspace_bytes = ( + q_latent_torch.numel() * q_latent_torch.element_size() + + q_rope_torch.numel() * q_rope_torch.element_size() + + c_latent_torch.numel() * c_latent_torch.element_size() + + c_rope_torch.numel() * c_rope_torch.element_size() + + o_torch.numel() * o_torch.element_size() + + lse_torch.numel() * lse_torch.element_size() + + cache_seqs_torch.numel() * cache_seqs_torch.element_size() + ) + one_workspace_bytes += ( + page_table_torch.numel() * page_table_torch.element_size() + ) + if is_var_split_kv: + one_workspace_bytes += ( + block_split_kvs_torch.numel() * block_split_kvs_torch.element_size() + ) + if workspace_torch is not None: + one_workspace_bytes += ( + workspace_torch.numel() * workspace_torch.element_size() + ) + workspace_count = testing.get_workspace_count( + one_workspace_bytes, warmup_iterations, iterations + ) + + avg_time_us = testing.benchmark( + compiled_mla, + workspace_generator=generate_tensors, + workspace_count=workspace_count, + stream=stream, + warmup_iterations=warmup_iterations, + iterations=iterations, + ) + + return avg_time_us # Return execution time in microseconds + + +if __name__ == "__main__": + + def parse_comma_separated_ints(s: str) -> Tuple[int, ...]: + try: + return tuple(int(x.strip()) for x in s.split(",")) + except ValueError: + raise argparse.ArgumentTypeError( + "Invalid format. Expected comma-separated integers." + ) + + def parse_mma_tiler(s: str) -> Tuple[int, int, Tuple[int, int]]: + ret = parse_comma_separated_ints(s) + if len(ret) != 2: + raise argparse.ArgumentTypeError( + "Invalid format. Expected 2 comma-separated integers." + ) + return (ret[0], ret[1]) + + parser = argparse.ArgumentParser(description="Example of MLA on Blackwell.") + + parser.add_argument( + "--in_dtype", + type=cutlass.dtype, + default=cutlass.Float16, + help="Input data type", + ) + + parser.add_argument( + "--out_dtype", + type=cutlass.dtype, + default=cutlass.Float16, + help="Output data type", + ) + + parser.add_argument( + "--acc_dtype", + type=cutlass.dtype, + default=cutlass.Float32, + help="Accumulator data type", + ) + + parser.add_argument( + "--lse_dtype", + type=cutlass.dtype, + default=cutlass.Float32, + help="LSE data type", + ) + parser.add_argument( + "--mma_qk_tiler_mn", + type=parse_mma_tiler, + default=(128, 128), + help="MMA tile shape (H, K)", + ) + parser.add_argument( + "--mma_pv_tiler_mn", + type=parse_mma_tiler, + default=(128, 256), + help="MMA tile shape (H, D)", + ) + + parser.add_argument( + "--is_persistent", + action="store_true", + help="Is persistent", + ) + + parser.add_argument( + "--batch_size", + type=int, + default=1, + help="Batch size", + ) + + parser.add_argument( + "--seq_len_q", + type=int, + default=1, + help="Sequence length of Q", + ) + + parser.add_argument( + "--seq_len_k", + type=int, + default=128, + help="Sequence length of K/V", + ) + + parser.add_argument( + "--num_heads", + type=int, + default=128, + help="Number of heads of Q", + ) + + parser.add_argument( + "--latent_dim", + type=int, + default=512, + help="Latent dimension of Q/C", + ) + + parser.add_argument( + "--rope_dim", + type=int, + default=64, + help="Rope dimension of Q/C", + ) + + parser.add_argument( + "--is_var_seq", + action="store_true", + help="Use variable length of sequence length or not", + ) + + parser.add_argument( + "--is_var_split_kv", + action="store_true", + help="Use variable length of split kv or not", + ) + + parser.add_argument( + "--page_size", + type=int, + default=128, + help="Page size of page table", + ) + + parser.add_argument( + "--split_kv", + type=int, + default=-1, + help="Split KV setting", + ) + + parser.add_argument( + "--softmax_scale", + type=float, + default=0.0416, + help="Scaling factor to scale softmax", + ) + + parser.add_argument( + "--output_scale", + type=float, + default=1.0, + help="Scaling factor to scale output", + ) + + parser.add_argument( + "--skip_correction_threshold", + type=float, + default=0.0, + help="Skip correction threshold", + ) + + parser.add_argument( + "--tolerance", type=float, default=1e-02, help="Tolerance for validation" + ) + + parser.add_argument( + "--warmup_iterations", + type=int, + default=0, + help="Number of iterations for warmup", + ) + + parser.add_argument( + "--iterations", + type=int, + default=1, + help="Number of iterations after warmup", + ) + + parser.add_argument( + "--skip_ref_check", + action="store_true", + help="Skip reference check", + ) + + parser.add_argument( + "--use_cold_l2", + action="store_true", + help="Use cold L2 cache", + ) + + args = parser.parse_args() + + run( + args.batch_size, + args.seq_len_q, + args.seq_len_k, + args.num_heads, + args.latent_dim, + args.rope_dim, + args.in_dtype, + args.out_dtype, + args.acc_dtype, + args.lse_dtype, + args.mma_qk_tiler_mn, + args.mma_pv_tiler_mn, + args.split_kv, + args.is_persistent, + args.is_var_seq, + args.is_var_split_kv, + args.page_size, + args.softmax_scale, + args.output_scale, + args.skip_correction_threshold, + args.tolerance, + args.warmup_iterations, + args.iterations, + args.skip_ref_check, + args.use_cold_l2, + ) + + print("PASS") diff --git a/examples/python/CuTeDSL/blackwell/mla/mla_decode_fp8.py b/examples/python/CuTeDSL/blackwell/mla/mla_decode_fp8.py new file mode 100644 index 00000000..e6383ef8 --- /dev/null +++ b/examples/python/CuTeDSL/blackwell/mla/mla_decode_fp8.py @@ -0,0 +1,4341 @@ +# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +import os +import sys +import argparse +import math +from typing import Type, Tuple, Optional +from types import SimpleNamespace + +import cuda.bindings.driver as cuda + +import cutlass +import cutlass.cute as cute +import cutlass.cute.testing as testing +from cutlass.cute.nvgpu import tcgen05 +from cutlass.cute.nvgpu.tcgen05 import OperandMajorMode +import cutlass.cute.nvgpu.cpasync as cpasync +import cutlass.utils as utils +import cutlass.pipeline as pipeline +from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait +import cutlass.utils.blackwell_helpers as sm100_utils +from cutlass.cute.runtime import from_dlpack +from cutlass.cute.arch import Arch +from cutlass.cutlass_dsl import BaseDSL + +if __name__ == "__main__": + current_dir = os.path.dirname(os.path.abspath(__file__)) + sys.path.insert(0, os.path.join(current_dir, "../..")) + +from blackwell.mla.mla_helpers import ( + ceil_div, + MAX_SPLITS, + LOG2_E, + MLAStaticTileScheduler, + MLAStaticTileSchedulerParams, + create_mla_static_tile_scheduler, + create_mla_static_tile_scheduler_params, +) + +""" +A Multi-Head Latent Attention (MLA) example using fp8 as input/output for the NVIDIA Blackwell SM100 architecture using CUTE DSL + +This example demonstrates an implementation of inference of multi-head latent attention using a TMA + Blackwell +SM100 TensorCore warp-specialized persistent kernel. The implementation integrates the (Qc + Qr)*(Kc + Kr)^T +matrix multiplication, softmax normalization, and softmax((Qc + Qr)*(Kc + Kr)^T)*Vc into a single kernel. +The kernel provides support for page table storage and variable-length KV cache sequences. It implements KV splitting +functionality to minimize latency when processing long KV sequences. + +The kernel implements key optimizations including: +- Warp specialization for different computation phases (load, MMA, softmax, correction, epilogue) +- Pipeline stages between different warps for overlapping computation and memory access +- Support for different precision data types +- Two sub-kernels (split KV kernel and reduction kernel) that enable split KV processing + +To run this example: + +.. code-block:: bash + + python examples/blackwell/mla_fp8.py \ + --batch_size 4 --latent_dim 512 --rope_dim 64 \ + --num_heads 128 --seq_len_q 1 --seq_len_k 1024 \ + --in_dtype Float8E4M3FN --out_dtype Float8E4M3FN \ + --acc_dtype Float32 --lse_dtype Float32 \ + --is_var_seq --is_var_split_kv \ + --is_persistent + +The above example runs Multi-Head Latent Attention (MLA) with the following configuration: +- Batch size: 4 +- Sequence length of Q: 1 +- Sequence length of K: 1024 +- Latent dimension: 512 +- RoPE dimension: 64 +- Number of heads: 128 +- Data types: Float8E4M3FN (input), Float8E4M3FN (output), Float32 (accumulation and LSE) + +It utilizes page table storage for the KV cache and enables both variable-length KV cache sequences +and variable split KV processing with persistent scheduling. + +To collect performance with NCU profiler: + +.. code-block:: bash + + ncu python examples/blackwell/mla_fp8.py \ + --batch_size 4 --latent_dim 512 --rope_dim 64 \ + --num_heads 128 --seq_len_q 1 --seq_len_k 1024 \ + --in_dtype Float8E4M3FN --out_dtype Float8E4M3FN \ + --acc_dtype Float32 --lse_dtype Float32 \ + --is_var_seq --is_var_split_kv \ + --is_persistent --warmup_iterations 3 \ + --iterations 10 --skip_ref_check + +Constraints for this example: +* Data type requirements: + - Input/output: Float8E4M3FN + - Accumulation and LSE: Float32 +* Fixed architecture parameters: + - Number of attention heads: 128 + - Latent dimension: 512 + - RoPE dimension: 64 +* Input query modes should be (NumHeads, LatentDim/RopeDim, SeqLenQ, BatchSize) +* Input kv latent/rope modes should be (SeqLenK, LatentDim/RopeDim, BatchSize) +* Query sequence length must be 1-4 +* Only supports 2-CTA instructions +* Variable sequence length requires page table storage enabled +""" + + +class BlackwellMultiHeadLatentAttentionForwardFP8: + def __init__( + self, + acc_dtype: Type[cutlass.Numeric], + lse_dtype: Type[cutlass.Numeric], + mma_qk_tiler_mn: Tuple[int, int], + mma_pv_tiler_mn: Tuple[int, int], + max_active_clusters: int, + page_size: int, + skip_correction_threshold: float, + is_persistent: bool, + is_var_seq: bool, + is_var_split_kv: bool, + ): + """Initializes the configuration for a Blackwell Multi-Head Latent Attention (MLA) kernel. + + :param acc_dtype: Data type for accumulation S and O + :type acc_dtype: Type[cutlass.Numeric] + :param lse_dtype: Data type for output LSE + :type lse_dtype: Type[cutlass.Numeric] + :param mma_s_tiler: The (H, K) tile shape of the MMA instruction for S + :type mma_s_tiler: Tuple[int, int] + :param mma_p_tiler: The (H, D) tile shape of the MMA instruction for P + :type mma_p_tiler: Tuple[int, int] + :param max_active_clusters: Maximum number of active clusters + :type max_active_clusters: int + :param page_size: The page size + :type page_size: int + :param skip_correction_threshold: Threshold to skip correction + :type skip_correction_threshold: float + :param is_persistent: Whether to use persistent kernel mode + :type is_persistent: bool + :param is_var_seq: Whether to use variable sequence length + :type is_var_seq: bool + :param is_var_split_kv: Whether to use variable split KV + :type is_var_split_kv: bool + """ + + self.latent_dim = 512 + self.rope_dim = 64 + self.acc_dtype = acc_dtype + self.lse_dtype = lse_dtype + self.mma_qk_tiler_mn = mma_qk_tiler_mn + self.mma_pv_tiler_mn = mma_pv_tiler_mn + self.max_active_clusters = max_active_clusters + self.skip_correction_threshold = skip_correction_threshold + self.is_persistent = is_persistent + self.page_size = page_size + self.is_var_seq = is_var_seq + self.is_var_split_kv = is_var_split_kv + self.cluster_shape_mnk = (2, 1, 1) + self.use_2cta_instrs = True + # When using 2 CTAs with m=128: warps 0-1 handle accumulation for first half [0, n/2), + # while warps 2-3 handle accumulation for second half [n/2, n) + self.warps_in_n = 2 + self.num_compute_warps = 4 + self.threads_per_warp = 32 + mma_qk_tiler_k = self.rope_dim * 2 + self.mma_qk_tiler = ( + self.mma_qk_tiler_mn[0], + self.mma_qk_tiler_mn[1], + mma_qk_tiler_k, + ) + self.mma_qk_rope_tiler = ( + self.mma_qk_tiler_mn[0], + self.mma_qk_tiler_mn[1], + self.rope_dim, + ) + self.mma_pv_tiler = ( + self.mma_pv_tiler_mn[0], + self.mma_pv_tiler_mn[1], + self.mma_qk_tiler[1] * self.mma_qk_tiler[2] // self.mma_pv_tiler_mn[1], + ) + self.iterations_qk_latent = self.latent_dim // self.mma_qk_tiler[2] + self.iterations_qk_rope = 1 + self.iterations_qk = self.iterations_qk_latent + self.iterations_qk_rope + self.iterations_pv_k = self.mma_qk_tiler[1] // self.mma_pv_tiler[2] + self.iterations_pv_n = self.latent_dim // self.mma_pv_tiler[1] + + # Set specialized warp ids + self.compute_warp_ids = (0, 1, 2, 3) + self.correction_warp_ids = (4, 5, 6, 7) + self.mma_warp_id = 8 + self.load_tma_k_warp_id = 9 + self.load_tma_v_warp_id = 10 + self.empty_warp_ids = (11,) + self.threads_per_cta = self.threads_per_warp * len( + ( + self.mma_warp_id, + self.load_tma_k_warp_id, + self.load_tma_v_warp_id, + *self.compute_warp_ids, + *self.correction_warp_ids, + *self.empty_warp_ids, + ) + ) + + # register settings + self.softmax_reg_num = 192 + self.correction_reg_num = 256 + self.other_reg_num = 48 + # Named barriers + self.tmem_ptr_sync_bar = pipeline.NamedBarrier( + barrier_id=1, + num_threads=( + self.threads_per_warp + + self.threads_per_warp * self.num_compute_warps * 2 + ), + ) + self.softmax_exchange_sync_bar = pipeline.NamedBarrier( + barrier_id=2, num_threads=(self.threads_per_warp * self.num_compute_warps) + ) + self.epilogue_exchange_sync_bar = pipeline.NamedBarrier( + barrier_id=3, num_threads=(self.threads_per_warp * self.num_compute_warps) + ) + + def _setup_attributes(self): + """Set up configurations and parameters for the MLA kernel operation. + + This method initializes and configures various attributes required for the + execution of the multi-head latent attention kernel, mainly about the pipeline stages: + + - Sets up staging parameters for Q, K, V inputs and accumulator data + - Configures pipeline stages for softmax, correction, and epilogue operations + """ + + self.load_q_stage = 1 + self.load_k_stage = 3 + self.load_v_stage = 2 + self.mma_s_stage = 2 + self.p_mma_stage = 2 + self.p_cor_stage = 2 + self.mma_o_stage = 2 + + self.tmem_o_offset = self.mma_s_stage * self.mma_qk_tiler[1] // self.warps_in_n + self.correction_factor_offset = ( + self.tmem_o_offset + self.latent_dim // self.warps_in_n + ) + + @cute.jit + def __call__( + self, + q_latent: cute.Tensor, + q_rope: cute.Tensor, + c_latent: cute.Tensor, + c_rope: cute.Tensor, + page_table: cute.Tensor, + o: cute.Tensor, + lse: cute.Tensor, + workspace: cute.Tensor, + split_kv: cutlass.Int32, + cache_seqs: Optional[cute.Tensor], + block_split_kvs: Optional[cute.Tensor], + softmax_scale: cutlass.Float32, + output_scale: cutlass.Float32, + stream: cuda.CUstream, + ): + """Execute the Multi-Head Latent Attention operation on the provided tensors. + + The method handles: + 1. Initialization of workspace for temporary split KV buffers + 2. Validation of tensor data types + 3. Initialization of hardware-specific parameters and memory layouts + 4. Configuration of TMA (Tensor Memory Access) operations + 5. Grid and work scheduling computation + 6. Kernel launch(split KV kernel and reduction kernel) with appropriate parameters + + :param q_latent: The query tensor with shape [num_head, latent_dim, seq_len_q, batch_size] + :type q_latent: cute.Tensor + :param q_rope: The query RoPE tensor with shape [num_head, rope_dim, seq_len_q, batch_size] + :type q_rope: cute.Tensor + :param c_latent: The key tensor with shape [seq_len_k, latent_dim, batch_size] + :type c_latent: cute.Tensor + :param c_rope: The key RoPE tensor with shape [seq_len_k, rope_dim, batch_size] + :type c_rope: cute.Tensor + :param page_table: The page table tensor with shape [page_count, batch_size] + :type page_table: cute.Tensor + :param o: The output tensor with shape [num_head, latent_dim, seq_len_q, batch_size] + :type o: cute.Tensor + :param lse: The LSE tensor with shape [num_head, seq_len_q, batch_size] + :type lse: cute.Tensor + :param workspace: The workspace tensor with 1-d shape prepared for acc_o and acc_lse + :type workspace: cute.Tensor + :param split_kv: The scalar factor for split KV + :type split_kv: cutlass.Int32 + :param cache_seqs: The cache sequences tensor with shape [batch_size] + :type cache_seqs: cute.Tensor + :param block_split_kvs: The block split KV tensor with shape [batch_size] + :type block_split_kvs: cute.Tensor + :param softmax_scale: The scale factor for softmax + :type softmax_scale: cutlass.Float32 + :param output_scale: The scale factor for the output + :type output_scale: cutlass.Float32 + :param stream: The CUDA stream to execute the kernel on + :type stream: cuda.CUstream + + :raises TypeError: If tensor data types don't match or aren't supported + """ + + # setup static attributes before smem/grid/tma computation + self.q_dtype = q_latent.element_type + self.k_dtype = c_latent.element_type + self.v_dtype = c_latent.element_type + self.o_dtype = o.element_type + + # check type consistency + if cutlass.const_expr( + self.q_dtype != self.k_dtype or self.q_dtype != self.v_dtype + ): + raise TypeError( + f"Type mismatch: {self.q_dtype} != {self.k_dtype} or {self.q_dtype} != {self.v_dtype}" + ) + # check leading dimensions of input/output + if cutlass.const_expr(q_latent.stride[1] != 1 or q_rope.stride[1] != 1): + raise ValueError("q_latent and q_rope must have leading dimension 1") + if cutlass.const_expr(c_latent.stride[1] != 1 or c_rope.stride[1] != 1): + raise ValueError("c_latent and c_rope must have leading dimension 1") + if cutlass.const_expr(o.stride[1] != 1): + raise ValueError("o must have leading dimension 1") + if cutlass.const_expr(lse.stride[0] != 1): + raise ValueError("lse must have leading dimension 0") + + acc_o, acc_lse = self.initialize_workspace( + q_latent.shape[0], + q_latent.shape[1], + q_latent.shape[2], + q_latent.shape[3], + split_kv, + self.acc_dtype, + workspace, + ) + + c_latent_tranpose_layout = cute.select(c_latent.layout, mode=[1, 0, 2]) + c_latent_transpose = cute.make_tensor( + c_latent.iterator, c_latent_tranpose_layout + ) + + self.q_major_mode = OperandMajorMode.K + self.k_major_mode = OperandMajorMode.K + self.v_major_mode = OperandMajorMode.MN + + self._setup_attributes() + + cta_group = tcgen05.CtaGroup.TWO + # the intermediate tensor p is from smem & k-major + p_major_mode = OperandMajorMode.K + qk_tiled_mma = sm100_utils.make_trivial_tiled_mma( + self.q_dtype, + self.q_major_mode, + self.k_major_mode, + self.acc_dtype, + cta_group, + self.mma_qk_tiler[:2], + ) + pv_tiled_mma = sm100_utils.make_trivial_tiled_mma( + self.v_dtype, + p_major_mode, + self.v_major_mode, + self.acc_dtype, + cta_group, + self.mma_pv_tiler[:2], + ) + + cta_layout_vmnk = cute.tiled_divide( + cute.make_layout(self.cluster_shape_mnk), + (qk_tiled_mma.thr_id.shape,), + ) + + self.epi_tile = self.mma_pv_tiler[:2] + + q_latent_smem_layout_staged = sm100_utils.make_smem_layout_a( + qk_tiled_mma, + self.mma_qk_tiler, + self.q_dtype, + (self.iterations_qk_latent * self.load_q_stage), + ) + q_latent_smem_layout_staged = cute.logical_divide( + q_latent_smem_layout_staged, (None, None, None, self.iterations_qk_latent) + ) + q_rope_smem_layout_staged = sm100_utils.make_smem_layout_a( + qk_tiled_mma, + self.mma_qk_rope_tiler, + self.q_dtype, + self.load_q_stage, + ) + + kc_latent_smem_layout_staged = sm100_utils.make_smem_layout_b( + qk_tiled_mma, + self.mma_qk_tiler, + self.k_dtype, + (self.iterations_qk_latent * self.load_k_stage), + ) + kc_page_tile_size = min( + self.page_size, qk_tiled_mma.op.shape_mnk[0] // qk_tiled_mma.thr_id.shape + ) + kc_latent_smem_layout_staged = cute.logical_divide( + kc_latent_smem_layout_staged, (None, None, None, self.iterations_qk_latent) + ) + + kc_latent_smem_layout_for_tma = sm100_utils.make_smem_layout( + OperandMajorMode.K, + (self.mma_qk_tiler[0] // qk_tiled_mma.thr_id.shape, self.mma_qk_tiler[2]), + self.k_dtype, + (self.iterations_qk_latent * self.load_k_stage), + ) + kc_latent_smem_layout_for_tma = cute.tiled_divide( + kc_latent_smem_layout_for_tma, (kc_page_tile_size, self.mma_qk_tiler[2]) + ) + kc_latent_smem_layout_for_tma = cute.logical_divide( + kc_latent_smem_layout_for_tma, (None, None, None, self.iterations_qk_latent) + ) + + kc_rope_smem_layout_staged = sm100_utils.make_smem_layout_b( + qk_tiled_mma, + self.mma_qk_rope_tiler, + self.k_dtype, + self.load_k_stage, + ) + kc_rope_smem_layout_for_tma = sm100_utils.make_smem_layout( + OperandMajorMode.K, + ( + self.mma_qk_rope_tiler[0] // qk_tiled_mma.thr_id.shape, + self.mma_qk_rope_tiler[2], + ), + self.k_dtype, + (self.iterations_qk_rope * self.load_k_stage), + ) + kc_rope_smem_layout_for_tma = cute.tiled_divide( + kc_rope_smem_layout_for_tma, (kc_page_tile_size, self.mma_qk_rope_tiler[2]) + ) + + p_smem_layout_staged = sm100_utils.make_smem_layout_a( + pv_tiled_mma, + self.mma_pv_tiler, + self.q_dtype, + (self.iterations_pv_k * self.p_mma_stage), + ) + p_smem_layout_staged = cute.logical_divide( + p_smem_layout_staged, (None, None, None, self.iterations_pv_k) + ) + + vc_smem_layout_staged = sm100_utils.make_smem_layout_b( + pv_tiled_mma, + self.mma_pv_tiler, + self.v_dtype, + (self.iterations_pv_k * self.iterations_pv_n * self.load_v_stage), + ) + vc_smem_layout_staged = cute.logical_divide( + cute.logical_divide( + vc_smem_layout_staged, + (None, None, None, self.iterations_pv_k * self.iterations_pv_n), + ), + (None, None, None, (self.iterations_pv_n, None)), + ) + vc_page_tile_size = min(self.page_size, self.mma_pv_tiler[2]) + vc_smem_layout_for_tma = sm100_utils.make_smem_layout( + OperandMajorMode.MN, + (self.mma_pv_tiler[1] // pv_tiled_mma.thr_id.shape, self.mma_pv_tiler[2]), + self.v_dtype, + (self.iterations_pv_k * self.iterations_pv_n * self.load_v_stage), + ) + vc_smem_layout_for_tma = cute.tiled_divide( + vc_smem_layout_for_tma, + ( + pv_tiled_mma.op.shape_mnk[1] // pv_tiled_mma.thr_id.shape, + vc_page_tile_size, + ), + ) + vc_smem_layout_for_tma = cute.logical_divide( + cute.logical_divide( + vc_smem_layout_for_tma, + (None, None, None, self.iterations_pv_k * self.iterations_pv_n), + ), + (None, None, None, (self.iterations_pv_n, None)), + ) + # TMA load for Q latent and rope + tma_load_op = cute.nvgpu.cpasync.CopyBulkTensorTileG2SOp(cta_group) + + q_smem_layout = cute.select(q_latent_smem_layout_staged, mode=[0, 1, 2]) + + tma_atom_q_latent, tma_tensor_q_latent = cute.nvgpu.make_tiled_tma_atom_A( + tma_load_op, + q_latent, + q_smem_layout, + self.mma_qk_tiler, + qk_tiled_mma, + cta_layout_vmnk.shape, + ) + q_rope_smem_layout = cute.select(q_rope_smem_layout_staged, mode=[0, 1, 2]) + tma_atom_q_rope, tma_tensor_q_rope = cute.nvgpu.make_tiled_tma_atom_A( + tma_load_op, + q_rope, + q_rope_smem_layout, + self.mma_qk_rope_tiler, + qk_tiled_mma, + cta_layout_vmnk.shape, + ) + # TMA load for c latent and k rope + kc_smem_layout = cute.select(kc_latent_smem_layout_for_tma, mode=[0]) + tma_atom_c_latent, tma_tensor_c_latent = self.make_paged_tiled_tma_atom( + tma_load_op, + c_latent, + kc_smem_layout, + (self.mma_qk_tiler[1], self.mma_qk_tiler[2]), + qk_tiled_mma, + is_k_load=True, + ) + kc_rope_smem_layout = cute.select(kc_rope_smem_layout_for_tma, mode=[0]) + tma_atom_c_rope, tma_tensor_c_rope = self.make_paged_tiled_tma_atom( + tma_load_op, + c_rope, + kc_rope_smem_layout, + (self.mma_qk_rope_tiler[1], self.mma_qk_rope_tiler[2]), + qk_tiled_mma, + is_k_load=True, + ) + + # TMA load for c latent transpose + vc_smem_layout = cute.select(vc_smem_layout_for_tma, mode=[0]) + tma_atom_c_latent_transpose, tma_tensor_c_latent_transpose = ( + self.make_paged_tiled_tma_atom( + tma_load_op, + c_latent_transpose, + vc_smem_layout, + (self.mma_pv_tiler[1], self.mma_pv_tiler[2]), + pv_tiled_mma, + is_k_load=False, + ) + ) + + q_latent_copy_size = ( + cute.size_in_bytes(self.q_dtype, q_smem_layout) + * cute.size(qk_tiled_mma.thr_id.shape) + * self.iterations_qk_latent + ) + q_rope_copy_size = ( + cute.size_in_bytes(self.q_dtype, q_rope_smem_layout) + * cute.size(qk_tiled_mma.thr_id.shape) + * self.iterations_qk_rope + ) + kc_latent_copy_size = ( + cute.size_in_bytes( + self.k_dtype, + cute.select(kc_latent_smem_layout_staged, mode=[0, 1, 2]), + ) + * cute.size(qk_tiled_mma.thr_id.shape) + * self.iterations_qk_latent + ) + kc_rope_copy_size = ( + cute.size_in_bytes( + self.k_dtype, + cute.select(kc_rope_smem_layout_staged, mode=[0, 1, 2]), + ) + * cute.size(qk_tiled_mma.thr_id.shape) + * self.iterations_qk_rope + ) + vc_copy_size = ( + cute.size_in_bytes( + self.v_dtype, cute.select(vc_smem_layout_staged, mode=[0, 1, 2]) + ) + * cute.size(pv_tiled_mma.thr_id.shape) + * self.iterations_pv_n + * self.iterations_pv_k + ) + + self.tma_copy_q_bytes = q_latent_copy_size + q_rope_copy_size + self.tma_copy_kc_bytes = kc_latent_copy_size + kc_rope_copy_size + self.tma_copy_vc_bytes = vc_copy_size + + tile_sched_params, grid = self._compute_grid( + o, + split_kv, + self.cluster_shape_mnk, + self.max_active_clusters, + self.is_persistent, + ) + + @cute.struct + class SplitKVKernelSharedStorage: + # Pipeline barriers + load_q_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.load_q_stage * 2] + load_k_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.load_k_stage * 2] + load_v_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.load_v_stage * 2] + mma_s_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.mma_s_stage * 2] + p_mma_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.p_mma_stage * 2] + p_cor_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.p_cor_stage * 2] + mma_o_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.mma_o_stage * 2] + + # Smem tensors + smem_p: cute.struct.Align[ + cute.struct.MemRange[self.q_dtype, cute.cosize(p_smem_layout_staged)], + 1024, + ] + smem_kc_latent: cute.struct.Align[ + cute.struct.MemRange[ + self.k_dtype, cute.cosize(kc_latent_smem_layout_staged) + ], + 1024, + ] + + smem_kc_rope: cute.struct.Align[ + cute.struct.MemRange[ + self.k_dtype, cute.cosize(kc_rope_smem_layout_staged) + ], + 1024, + ] + smem_q_latent: cute.struct.Align[ + cute.struct.MemRange[ + self.q_dtype, cute.cosize(q_latent_smem_layout_staged) + ], + 1024, + ] + smem_q_rope: cute.struct.Align[ + cute.struct.MemRange[ + self.q_dtype, cute.cosize(q_rope_smem_layout_staged) + ], + 1024, + ] + smem_vc: cute.struct.Align[ + cute.struct.MemRange[self.v_dtype, cute.cosize(vc_smem_layout_staged)], + 1024, + ] + softmax_smem_exchange: cute.struct.MemRange[ + self.acc_dtype, self.num_compute_warps * self.threads_per_warp + ] + epilogue_smem_exchange: cute.struct.MemRange[ + self.acc_dtype, self.num_compute_warps * self.threads_per_warp + ] + + # Tmem dealloc cluster barrier + tmem_dealloc_mbar_ptr: cutlass.Int64 + + # Tmem holding buffer + tmem_holding_buf: cutlass.Int32 + + softmax_scale_log2 = softmax_scale * LOG2_E + + self.split_kv_kernel( + qk_tiled_mma, + pv_tiled_mma, + tma_atom_q_latent, + tma_tensor_q_latent, + tma_atom_q_rope, + tma_tensor_q_rope, + tma_atom_c_latent, + tma_tensor_c_latent, + tma_atom_c_rope, + tma_tensor_c_rope, + tma_atom_c_latent_transpose, + tma_tensor_c_latent_transpose, + page_table, + o, + lse, + acc_o, + acc_lse, + split_kv, + cache_seqs, + block_split_kvs, + softmax_scale_log2, + output_scale, + q_latent_smem_layout_staged, + q_rope_smem_layout_staged, + kc_latent_smem_layout_staged, + kc_rope_smem_layout_staged, + p_smem_layout_staged, + vc_smem_layout_staged, + kc_latent_smem_layout_for_tma, + kc_rope_smem_layout_for_tma, + vc_smem_layout_for_tma, + cta_layout_vmnk, + tile_sched_params, + SplitKVKernelSharedStorage, + ).launch( + grid=grid, + block=[self.threads_per_cta, 1, 1], + cluster=self.cluster_shape_mnk, + smem=SplitKVKernelSharedStorage.size_in_bytes(), + stream=stream, + min_blocks_per_mp=1, + ) + if cutlass.const_expr(acc_o is not None): + self.reduction_kernel( + o, + lse, + acc_o, + acc_lse, + split_kv, + cache_seqs, + block_split_kvs, + ).launch( + grid=(q_latent.shape[0], q_latent.shape[2], q_latent.shape[3]), + block=[self.threads_per_warp * self.num_compute_warps, 1, 1], + smem=MAX_SPLITS * self.acc_dtype.width // 8, + stream=stream, + min_blocks_per_mp=1, + ) + + @cute.jit + def make_paged_tiled_tma_atom( + self, + tma_load_op: cute.nvgpu.cpasync.CopyBulkTensorTileG2SOp, + gmem: cute.Tensor, + smem_layout: cute.Layout, + mma_tiler, + tiled_mma: cute.TiledMma, + is_k_load: bool, + ): + ident = cute.make_identity_layout(gmem.shape) + g_tile = cute.composition(ident, mma_tiler) + cta_mn = mma_tiler[0] // tiled_mma.thr_id.shape + cta_v_map = cute.flat_divide(g_tile, (cta_mn,)) + cta_v_map = cute.select(cta_v_map, mode=[0, 2]) + page_tile_size = ( + min(self.page_size, cta_mn) + if is_k_load + else min(self.page_size, mma_tiler[1]) + ) + cta_v_map = cute.zipped_divide( + cta_v_map, + (page_tile_size, mma_tiler[1]) if is_k_load else (cta_mn, page_tile_size), + ) + cta_v_map = cute.select(cta_v_map, mode=[0]) + from cutlass._mlir.dialects import cute_nvgpu as _cute_nvgpu_ir + + res = _cute_nvgpu_ir.atom_make_non_exec_tiled_tma_load( + gmem.value, + smem_layout.value, + cta_v_map, + tma_load_op._to_ir(), + num_multicast=1, + ) + return ( + cute.CopyAtom( + tma_load_op, cpasync.CopyBulkTensorTileG2SNonExecTrait(res[0]) + ), + res[1], + ) + + @cute.kernel + def split_kv_kernel( + self, + tiled_mma_qk: cute.TiledMma, + tiled_mma_pv: cute.TiledMma, + tma_atom_q_latent: Optional[cute.CopyAtom], + mQL: cute.Tensor, + tma_atom_q_rope: Optional[cute.CopyAtom], + mQR: cute.Tensor, + tma_atom_c_latent: Optional[cute.CopyAtom], + mCL: cute.Tensor, + tma_atom_c_rope: Optional[cute.CopyAtom], + mKR: cute.Tensor, + tma_atom_c_latent_transpose: Optional[cute.CopyAtom], + mCLT: cute.Tensor, + mPT: cute.Tensor, + mO: Optional[cute.Tensor], + mLSE: Optional[cute.Tensor], + mAccO: Optional[cute.Tensor], + mAccLSE: Optional[cute.Tensor], + split_kv: cutlass.Int32, + cache_seqs: cute.Tensor, + block_split_kvs: cute.Tensor, + softmax_scale_log2: cutlass.Float32, + output_scale: cutlass.Float32, + q_latent_smem_layout_staged: cute.ComposedLayout, + q_rope_smem_layout_staged: cute.ComposedLayout, + kc_latent_smem_layout_staged: cute.ComposedLayout, + kc_rope_smem_layout_staged: cute.ComposedLayout, + p_smem_layout_staged: cute.ComposedLayout, + vc_smem_layout_staged: cute.ComposedLayout, + kc_latent_smem_layout_for_tma: Optional[cute.ComposedLayout], + kc_rope_smem_layout_for_tma: Optional[cute.ComposedLayout], + vc_smem_layout_for_tma: Optional[cute.ComposedLayout], + cta_layout_vmnk: cute.Layout, + tile_sched_params: MLAStaticTileSchedulerParams, + SharedStorage: cutlass.Constexpr, + ): + """The device split_kv kernel implementation of the Multi-Head Latent Attention. + + This kernel coordinates multiple specialized warps to perform different phases of the MLA computation: + 1. Load warp: Loads Q/C latent/rope data from global memory to shared memory using TMA + 2. MMA warp: Performs matrix multiplications (Q*K^T and P*V) + 3. Compute warps: Compute softmax and do rescaling on accumulators, and store the intermediate/final results + to global memory + + The kernel produces either intermediate or final results of the MLA computation based on the split_kv parameter. + When split_kv is 1, the kernel generates the final results directly. Otherwise, it produces intermediate results + that will later be combined by a reduction kernel. + + The kernel implements a complex pipeline with overlapping computation and memory operations, + using tensor memory access (TMA) for efficient data loading, warp specialization for different + computation phases. + + :param tiled_mma_qk: Tiled MMA for Q*K^T + :type tiled_mma_qk: cute.TiledMma + :param tiled_mma_pv: Tiled MMA for P*V + :type tiled_mma_pv: cute.TiledMma + :param tma_atom_q_latent: TMA copy atom for query latent tensor + :type tma_atom_q_latent: cute.CopyAtom + :param mQL: query latent tensor + :type mQL: cute.Tensor + :param tma_atom_q_rope: TMA copy atom for query rope tensor + :type tma_atom_q_rope: cute.CopyAtom + :param mKR: Compressed rope tensor + :type mKR: cute.Tensor + :param tma_atom_c_latent: TMA copy atom for c latent tensor + :type tma_atom_c_latent: cute.CopyAtom + :param mCL: Compressed latent tensor + :type mCL: cute.Tensor + :param tma_atom_c_rope: TMA copy atom for c rope tensor + :type tma_atom_c_rope: cute.CopyAtom + :param mCLT: Compressed latent transpose tensor + :type mCLT: cute.Tensor + :param mPT: Page table tensor + :type mPT: cute.Tensor + :param mO: Output tensor + :type mO: cute.Tensor + :param mLSE: Log-sum-exp tensor + :type mLSE: cute.Tensor + :param mAccO: Intermediate accumulator output tensor + :type mAccO: cute.Tensor + :param mAccLSE: Intermediate accumulator log-sum-exp tensor + :type mAccLSE: cute.Tensor + :param split_kv: The split_kv parameter + :type split_kv: cutlass.Int32 + :param cache_seqs: The variable sequence length tensor + :type cache_seqs: cute.Tensor + :param block_split_kvs: The per-block split_kv values tensor + :type block_split_kvs: cute.Tensor + :param softmax_scale_log2: The log2 scale factor for softmax + :type softmax_scale_log2: cutlass.Float32 + :param output_scale: The scale factor for the output + :type output_scale: cutlass.Float32 + :param q_latent_smem_layout_staged: Shared memory layout for query tensor + :type q_latent_smem_layout_staged: cute.ComposedLayout + :param q_rope_smem_layout_staged: Shared memory layout for query rope tensor + :type q_rope_smem_layout_staged: cute.ComposedLayout + :param kc_latent_smem_layout_staged: Shared memory layout for key tensor + :type kc_latent_smem_layout_staged: cute.ComposedLayout + :param kc_rope_smem_layout_staged: Shared memory layout for key rope tensor + :type kc_rope_smem_layout_staged: cute.ComposedLayout + :param p_smem_layout_staged: Shared memory layout for probability matrix + :type p_smem_layout_staged: cute.ComposedLayout + :param vc_smem_layout_staged: Shared memory layout for value tensor + :type vc_smem_layout_staged: cute.ComposedLayout + :param cta_layout_vmnk: Layout for compute threads + :type cta_layout_vmnk: cute.Layout + :param tile_sched_params: Scheduling parameters for work distribution + :type tile_sched_params: MLAStaticTileSchedulerParams + :param SharedStorage: Shared storage for the kernel + :type SharedStorage: cutlass.Constexpr + """ + + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + + tidx, _, _ = cute.arch.thread_idx() + bidx, _, _ = cute.arch.block_idx() + mma_tile_coord_v = bidx % cute.size(tiled_mma_qk.thr_id.shape) + is_leader_cta = mma_tile_coord_v == 0 + + # Prefetch tma descriptor + if warp_idx == self.mma_warp_id: + cpasync.prefetch_descriptor(tma_atom_q_latent) + cpasync.prefetch_descriptor(tma_atom_q_rope) + cpasync.prefetch_descriptor(tma_atom_c_latent) + cpasync.prefetch_descriptor(tma_atom_c_rope) + cpasync.prefetch_descriptor(tma_atom_c_latent_transpose) + + # Alloc + smem = utils.SmemAllocator() + storage = smem.allocate(SharedStorage) + + # Tensor memory dealloc barrier init + tmem = utils.TmemAllocator( + storage.tmem_holding_buf, + barrier_for_retrieve=self.tmem_ptr_sync_bar, + allocator_warp_id=self.mma_warp_id, + is_two_cta=self.use_2cta_instrs, + two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar_ptr, + ) + + load_q_pipeline = self.make_and_init_load_qkv_pipeline( + storage.load_q_mbar_ptr.data_ptr(), + cta_layout_vmnk, + self.load_q_stage, + self.tma_copy_q_bytes, + ) + load_k_pipeline = self.make_and_init_load_qkv_pipeline( + storage.load_k_mbar_ptr.data_ptr(), + cta_layout_vmnk, + self.load_k_stage, + self.tma_copy_kc_bytes, + ) + load_v_pipeline = self.make_and_init_load_qkv_pipeline( + storage.load_v_mbar_ptr.data_ptr(), + cta_layout_vmnk, + self.load_v_stage, + self.tma_copy_vc_bytes, + ) + mma_s_pipeline = self.make_and_init_mma_s_pipeline( + storage.mma_s_mbar_ptr.data_ptr(), cta_layout_vmnk + ) + p_mma_pipeline = self.make_and_init_p_mma_pipeline( + storage.p_mma_mbar_ptr.data_ptr(), cta_layout_vmnk + ) + p_cor_pipeline = self.make_and_init_p_cor_pipeline( + storage.p_cor_mbar_ptr.data_ptr() + ) + mma_o_pipeline = self.make_and_init_mma_o_pipeline( + storage.mma_o_mbar_ptr.data_ptr(), cta_layout_vmnk + ) + + # Cluster arrive after barrier init + pipeline_init_arrive(cluster_shape_mn=self.cluster_shape_mnk, is_relaxed=True) + + # Generate smem tensor Q/KC/VC/exchange + # (MMA, MMA_H, MMA_R, PIPE) + sQ = storage.smem_q_latent.get_tensor( + q_latent_smem_layout_staged.outer, swizzle=q_latent_smem_layout_staged.inner + ) + sQ_rope = storage.smem_q_rope.get_tensor( + q_rope_smem_layout_staged.outer, swizzle=q_rope_smem_layout_staged.inner + ) + # (MMA, MMA_K, MMA_R, PIPE) + sKC = storage.smem_kc_latent.get_tensor( + kc_latent_smem_layout_staged.outer, + swizzle=kc_latent_smem_layout_staged.inner, + ) + sKC_rope = storage.smem_kc_rope.get_tensor( + kc_rope_smem_layout_staged.outer, swizzle=kc_rope_smem_layout_staged.inner + ) + sKC_for_tma = storage.smem_kc_latent.get_tensor( + kc_latent_smem_layout_for_tma.outer, + swizzle=kc_latent_smem_layout_for_tma.inner, + ) + sKC_rope_for_tma = storage.smem_kc_rope.get_tensor( + kc_rope_smem_layout_for_tma.outer, swizzle=kc_rope_smem_layout_for_tma.inner + ) + # (MMA, MMA_D, MMA_K, PIPE) + sVC = storage.smem_vc.get_tensor( + vc_smem_layout_staged.outer, swizzle=vc_smem_layout_staged.inner + ) + sVC_for_tma = storage.smem_vc.get_tensor( + vc_smem_layout_for_tma.outer, swizzle=vc_smem_layout_for_tma.inner + ) + # (MMA, MMA_H, MMA_K) + sP = storage.smem_p.get_tensor( + p_smem_layout_staged.outer, swizzle=p_smem_layout_staged.inner + ) + # (compute_threads,) + softmax_smem_exchange = storage.softmax_smem_exchange.get_tensor( + cute.make_layout(self.num_compute_warps * self.threads_per_warp) + ) + epilogue_smem_exchange = storage.epilogue_smem_exchange.get_tensor( + cute.make_layout(self.num_compute_warps * self.threads_per_warp) + ) + + # + # Cluster wait before tensor memory alloc + # + pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mnk) + + # /////////////////////////////////////////////////////////////////////////////// + # Load warps, including page table and data tensors + # /////////////////////////////////////////////////////////////////////////////// + if warp_idx >= self.empty_warp_ids[0] and warp_idx <= self.empty_warp_ids[-1]: + cute.arch.setmaxregister_decrease(self.other_reg_num) + + if warp_idx == self.load_tma_k_warp_id: + cute.arch.setmaxregister_decrease(self.other_reg_num) + load_q_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.load_q_stage + ) + load_k_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.load_k_stage + ) + tile_sched = create_mla_static_tile_scheduler( + tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim() + ) + work_tile = tile_sched.initial_work_tile_info() + while work_tile.is_valid_tile: + blk_coord = work_tile.tile_idx + k_index, k_tile_count, local_split_kv = self.get_k_tile_count( + split_kv, + cache_seqs, + block_split_kvs, + blk_coord, + ) + if k_tile_count > 0: + # Construct fixed common/tma_qk/tma_pv params for load_tma + tma_common_params = SimpleNamespace( + blk_coord=blk_coord, + local_split_kv=local_split_kv, + load_q_pipeline=load_q_pipeline, + load_k_pipeline=load_k_pipeline, + load_v_pipeline=load_v_pipeline, + mPT=mPT, + ) + tma_qk_params = SimpleNamespace( + tiled_mma_qk=tiled_mma_qk, + tma_atom_q_latent=tma_atom_q_latent, + tma_atom_q_rope=tma_atom_q_rope, + tma_atom_c_latent=tma_atom_c_latent, + tma_atom_c_rope=tma_atom_c_rope, + mQL=mQL, + mQR=mQR, + mCL=mCL, + mKR=mKR, + sQ=sQ, + sQ_rope=sQ_rope, + sKC=sKC_for_tma, + sKC_rope=sKC_rope_for_tma, + ) + # Load tma + load_q_producer_state, load_k_producer_state = self.load_tma_qk( + tma_common_params, + tma_qk_params, + k_index, + k_tile_count, + load_q_producer_state, + load_k_producer_state, + ) + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + + load_q_pipeline.producer_tail(load_q_producer_state) + load_k_pipeline.producer_tail(load_k_producer_state) + + if warp_idx == self.load_tma_v_warp_id: + cute.arch.setmaxregister_decrease(self.other_reg_num) + load_v_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.load_v_stage + ) + tile_sched = create_mla_static_tile_scheduler( + tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim() + ) + work_tile = tile_sched.initial_work_tile_info() + while work_tile.is_valid_tile: + blk_coord = work_tile.tile_idx + k_index, k_tile_count, local_split_kv = self.get_k_tile_count( + split_kv, + cache_seqs, + block_split_kvs, + blk_coord, + ) + if k_tile_count > 0: + # Construct fixed common/tma_qk/tma_pv params for load_tma + tma_common_params = SimpleNamespace( + blk_coord=blk_coord, + local_split_kv=local_split_kv, + load_v_pipeline=load_v_pipeline, + mPT=mPT, + ) + tma_pv_params = SimpleNamespace( + tiled_mma_pv=tiled_mma_pv, + tma_atom_c_latent_transpose=tma_atom_c_latent_transpose, + mCLT=mCLT, + sVC=sVC_for_tma, + ) + # Load tma + load_v_producer_state = self.load_tma_v( + tma_common_params, + tma_pv_params, + k_index, + k_tile_count, + load_v_producer_state, + ) + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + load_v_pipeline.producer_tail(load_v_producer_state) + + # /////////////////////////////////////////////////////////////////////////////// + # MMA warp + # /////////////////////////////////////////////////////////////////////////////// + if warp_idx == self.mma_warp_id: + cute.arch.setmaxregister_decrease(self.other_reg_num) + # Alloc tensor memory buffer + tmem.allocate(cute.arch.get_max_tmem_alloc_cols("sm_100")) + tmem.wait_for_alloc() + tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + + load_q_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.load_q_stage + ) + load_k_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.load_k_stage + ) + load_v_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.load_v_stage + ) + mma_s_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.mma_s_stage + ) + p_mma_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.p_mma_stage + ) + mma_o_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.mma_o_stage + ) + tile_sched = create_mla_static_tile_scheduler( + tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim() + ) + work_tile = tile_sched.initial_work_tile_info() + while work_tile.is_valid_tile: + blk_coord = work_tile.tile_idx + k_index, k_tile_count, local_split_kv = self.get_k_tile_count( + split_kv, cache_seqs, block_split_kvs, blk_coord + ) + if k_tile_count > 0: + mma_common_params = SimpleNamespace( + blk_coord=blk_coord, + local_split_kv=local_split_kv, + load_q_pipeline=load_q_pipeline, + load_k_pipeline=load_k_pipeline, + load_v_pipeline=load_v_pipeline, + tmem_ptr=tmem_ptr, + is_leader_cta=is_leader_cta, + L=mCL.shape[1], + ) + mma_qk_params = SimpleNamespace( + mma_s_pipeline=mma_s_pipeline, + sQ=sQ, + sQ_rope=sQ_rope, + sKC=sKC, + sKC_rope=sKC_rope, + ) + mma_pv_params = SimpleNamespace( + p_mma_pipeline=p_mma_pipeline, + mma_o_pipeline=mma_o_pipeline, + sP=sP, + sVC=sVC, + ) + ( + tiled_mma_qk, + tiled_mma_pv, + load_q_consumer_state, + load_k_consumer_state, + load_v_consumer_state, + mma_s_producer_state, + p_mma_consumer_state, + mma_o_producer_state, + ) = self.mma( + mma_common_params, + mma_qk_params, + mma_pv_params, + k_tile_count, + tiled_mma_qk, + tiled_mma_pv, + load_q_consumer_state, + load_k_consumer_state, + load_v_consumer_state, + mma_s_producer_state, + p_mma_consumer_state, + mma_o_producer_state, + ) + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + + mma_s_pipeline.producer_tail(mma_s_producer_state) + mma_o_pipeline.producer_tail(mma_o_producer_state) + + tmem.relinquish_alloc_permit() + tmem.free(tmem_ptr) + + # /////////////////////////////////////////////////////////////////////////////// + # Compute warp + # /////////////////////////////////////////////////////////////////////////////// + if ( + warp_idx >= self.compute_warp_ids[0] + and warp_idx <= self.compute_warp_ids[-1] + ): + cute.arch.setmaxregister_increase(self.softmax_reg_num) + mma_s_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.mma_s_stage + ) + p_mma_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.p_mma_stage + ) + p_cor_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.p_cor_stage + ) + mma_o_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.mma_o_stage + ) + tmem.wait_for_alloc() + tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + + tile_sched = create_mla_static_tile_scheduler( + tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim() + ) + work_tile = tile_sched.initial_work_tile_info() + while work_tile.is_valid_tile: + blk_coord = work_tile.tile_idx + k_index, k_tile_count, local_split_kv = self.get_k_tile_count( + split_kv, cache_seqs, block_split_kvs, blk_coord + ) + if k_tile_count > 0: + compute_common_params = SimpleNamespace( + blk_coord=blk_coord, + split_kv=split_kv, + local_split_kv=local_split_kv, + smem_exchange=softmax_smem_exchange, + mAccO=mAccO, + mO=mO, + K=cache_seqs[blk_coord[2]], + L=mCL.shape[1], + tmem_ptr=tmem_ptr, + tidx=tidx, + p_cor_pipeline=p_cor_pipeline, + ) + compute_softmax_params = SimpleNamespace( + tiled_mma_qk=tiled_mma_qk, + sP=sP, + mma_s_pipeline=mma_s_pipeline, + p_mma_pipeline=p_mma_pipeline, + softmax_scale_log2=softmax_scale_log2, + ) + mma_s_consumer_state, p_mma_producer_state, p_cor_producer_state = ( + self.compute( + compute_common_params, + compute_softmax_params, + k_index=k_index, + k_tile_count=k_tile_count, + mma_s_consumer_state=mma_s_consumer_state, + p_mma_producer_state=p_mma_producer_state, + p_cor_producer_state=p_cor_producer_state, + ) + ) + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + p_cor_pipeline.producer_tail(p_cor_producer_state) + + # /////////////////////////////////////////////////////////////////////////////// + # Correction warp + # /////////////////////////////////////////////////////////////////////////////// + if ( + warp_idx >= self.correction_warp_ids[0] + and warp_idx <= self.correction_warp_ids[-1] + ): + cute.arch.setmaxregister_increase(self.correction_reg_num) + p_cor_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.p_cor_stage + ) + mma_o_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.mma_o_stage + ) + # sync with mma warp before retrieving tmem ptr + tmem.wait_for_alloc() + + tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + + tile_sched = create_mla_static_tile_scheduler( + tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim() + ) + work_tile = tile_sched.initial_work_tile_info() + while work_tile.is_valid_tile: + blk_coord = work_tile.tile_idx + k_index, k_tile_count, local_split_kv = self.get_k_tile_count( + split_kv, cache_seqs, block_split_kvs, blk_coord + ) + if k_tile_count > 0: + compute_common_params = SimpleNamespace( + blk_coord=blk_coord, + split_kv=split_kv, + local_split_kv=local_split_kv, + smem_exchange=epilogue_smem_exchange, + mAccO=mAccO, + mO=mO, + K=cache_seqs[blk_coord[2]], + L=mCL.shape[1], + H=mQL.shape[0], + tmem_ptr=tmem_ptr, + tidx=tidx, + tiled_mma_pv=tiled_mma_pv, + p_cor_pipeline=p_cor_pipeline, + mma_o_pipeline=mma_o_pipeline, + ) + compute_epilogue_params = SimpleNamespace( + output_scale=output_scale, + softmax_scale_log2=softmax_scale_log2, + mAccLSE=mAccLSE, + mLSE=mLSE, + ) + p_cor_consumer_state, mma_o_consumer_state = self.correction( + compute_common_params, + compute_epilogue_params, + k_tile_count=k_tile_count, + p_cor_consumer_state=p_cor_consumer_state, + mma_o_consumer_state=mma_o_consumer_state, + ) + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + + return + + @cute.kernel + def reduction_kernel( + self, + mO: cute.Tensor, + mLSE: cute.Tensor, + mAccO: cute.Tensor, + mAccLSE: cute.Tensor, + split_kv: cutlass.Int32, + cache_seqs: cute.Tensor, + block_split_kvs: cute.Tensor, + ): + """The reduction kernel for Multi-Head Latent Attention (MLA) that combines intermediate results + from multiple split_kv blocks into final outputs. + + :param mO: Output tensor for storing final results + :type mO: cute.Tensor + :param mLSE: Log-sum-exp tensor for storing final LSE values + :type mLSE: cute.Tensor + :param mAccO: Accumulated output tensor from split_kv blocks + :type mAccO: cute.Tensor + :param mAccLSE: Accumulated LSE tensor from split_kv blocks + :type mAccLSE: cute.Tensor + :param split_kv: Number of split_kv blocks + :type split_kv: cutlass.Int32 + :param cache_seqs: Cache sequence lengths tensor + :type cache_seqs: cute.Tensor + :param block_split_kvs: Per-block split_kv values tensor (for variable split_kv) + :type block_split_kvs: cute.Tensor + """ + bidx, bidy, bidz = cute.arch.block_idx() + tidx, _, _ = cute.arch.thread_idx() + blk_coord = (bidx, bidy, bidz) + local_split_kv = ( + block_split_kvs[blk_coord[2]] if self.is_var_split_kv else split_kv + ) + k_tile_total = cute.ceil_div(cache_seqs[blk_coord[2]], self.mma_qk_tiler[1]) + k_tile_per_cta = cute.ceil_div(k_tile_total, local_split_kv) + local_split_kv = cute.ceil_div(k_tile_total, k_tile_per_cta) + + # Alloc shared memory + smem = utils.SmemAllocator() + storage = smem.allocate(MAX_SPLITS * self.acc_dtype.width // 8, 16) + lse_scale_ptr = cute.recast_ptr(storage, dtype=self.acc_dtype) + smem_lse_scale = cute.make_tensor(lse_scale_ptr, cute.make_layout(MAX_SPLITS)) + + gLSE = mAccLSE[blk_coord[0], None, blk_coord[1], blk_coord[2]] + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + if warp_idx == 0: + # calculate the global lse and exp ^ (local_lse - global_lse) + lse_per_thread = cute.ceil_div(MAX_SPLITS, self.threads_per_warp) + + local_lse = cute.make_rmem_tensor( + cute.make_layout(lse_per_thread), self.lse_dtype + ) + lse_max = -self.lse_dtype.inf + # find the max lse + for i in cutlass.range_constexpr(lse_per_thread): + split_kv_idx = tidx + i * self.threads_per_warp + local_lse[i] = ( + gLSE[split_kv_idx] + if cute.elem_less(split_kv_idx, local_split_kv) + else -self.lse_dtype.inf + ) + # reduce the local lse + lse_max = cute.arch.fmax(lse_max, local_lse[i]) + lse_max = cute.arch.warp_reduction_max(lse_max) + lse_max = lse_max if lse_max != -self.lse_dtype.inf else 0.0 + # calculate sum_lse + sum_lse = 0.0 + for i in cutlass.range_constexpr(lse_per_thread): + sum_lse += cute.math.exp2(local_lse[i] - lse_max, fastmath=True) + sum_lse = cute.arch.warp_reduction_sum(sum_lse) + # calculate the global_lse + global_lse = ( + lse_max + cute.math.log2(sum_lse, fastmath=True) + if not sum_lse == self.lse_dtype(0.0) or sum_lse != sum_lse + else self.lse_dtype.inf + ) + if tidx == 0: + mLSE[blk_coord[0], blk_coord[1], blk_coord[2]] = global_lse + # store the scale to shared memory + for i in cutlass.range_constexpr(lse_per_thread): + split_kv_idx = tidx + i * self.threads_per_warp + if cute.elem_less(split_kv_idx, local_split_kv): + smem_lse_scale[split_kv_idx] = cute.math.exp2( + local_lse[i] - global_lse, fastmath=True + ) + + pipeline.sync(barrier_id=4) + + elements_per_thread = cute.ceil_div( + self.latent_dim, self.threads_per_warp * self.num_compute_warps + ) + gAccO = mAccO[blk_coord[0], None, None, blk_coord[1], blk_coord[2]] + rAccO = cute.make_rmem_tensor( + cute.make_layout(elements_per_thread), self.acc_dtype + ) + rO = cute.make_rmem_tensor(cute.make_layout(elements_per_thread), self.o_dtype) + rAccO.fill(0.0) + for i in range(local_split_kv): + for j in cutlass.range_constexpr(elements_per_thread): + element_idx = tidx + j * self.threads_per_warp * self.num_compute_warps + rAccO[j] += gAccO[i, element_idx] * smem_lse_scale[i] + rO.store(rAccO.load().to(self.o_dtype)) + for j in cutlass.range_constexpr(elements_per_thread): + element_idx = tidx + j * self.threads_per_warp * self.num_compute_warps + mO[blk_coord[0], element_idx, blk_coord[1], blk_coord[2]] = rO[j] + return + + @staticmethod + def get_split_kv( + B: int, S: int, K: int, mma_qk_tiler_mn: tuple, max_active_blocks: int + ) -> int: + """Get the proper split_kv value for the MLA kernel based on parameters. + + :param B: Batch size + :type B: int + :param S: Sequence length + :type S: int + :param K: Sequence length + :type K: int + :param mma_qk_tiler_mn: MLA tiling parameters + :type mma_qk_tiler_mn: tuple + :param max_active_blocks: Maximum number of active blocks + :type max_active_blocks: int + :return: Split_kv value + :rtype: int + """ + max_splits = ceil_div(K, mma_qk_tiler_mn[1]) + blocks_per_batch = max(1, max_active_blocks // B // (S * 2)) + split_heur = min(max_splits, blocks_per_batch) + k_waves = ceil_div(max_splits, split_heur) + split_wave_aware = ceil_div(max_splits, k_waves) + max_split_kv = 32 + return min(split_wave_aware, max_split_kv) + + @cute.jit + def get_k_tile_count( + self, + split_kv: cutlass.Int32, + cache_seqs: cute.Tensor, + block_split_kvs: cute.Tensor, + blk_coord: cute.Coord, + ) -> tuple[cutlass.Int32, cutlass.Int32, cutlass.Int32]: + """Get the current k_index, k_tile_count, and local split_kv value for the MLA kernel. + + :param split_kv: Split_kv value + :type split_kv: cutlass.Int32 + :param cache_seqs: Cache sequence lengths tensor + :type cache_seqs: cute.Tensor + :param block_split_kvs: Per-block split_kv values tensor + :type block_split_kvs: cute.Tensor + :param blk_coord: Block coordinate + :type blk_coord: cute.Coord + :return: k_index, k_tile_count, split_kv + :rtype: tuple[cutlass.Int32, cutlass.Int32, cutlass.Int32] + """ + K = cache_seqs[blk_coord[2]] + if cutlass.const_expr(self.is_var_split_kv): + split_kv = block_split_kvs[blk_coord[2]] + + k_tile_total = cute.ceil_div(K, self.mma_qk_tiler[1]) + k_tile_per_cta = cute.ceil_div(k_tile_total, split_kv) + k_index = blk_coord[3] * k_tile_per_cta + k_tile_count = max(0, min(k_tile_total, k_index + k_tile_per_cta) - k_index) + return k_index, k_tile_count, split_kv + + @cute.jit + def load_tma_qk( + self, + common_params: SimpleNamespace, + qk_params: SimpleNamespace, + k_index: cutlass.Int32, + k_tile_count: cutlass.Int32, + load_q_producer_state: pipeline.PipelineState | None = None, + load_k_producer_state: pipeline.PipelineState | None = None, + ) -> tuple[pipeline.PipelineState, pipeline.PipelineState]: + """Load wrap to load Q/K tensors. Updates the load qk producer state. + + :param common_params: The common parameters + :type common_params: SimpleNamespace + :param qk_params: The qk parameters + :type qk_params: SimpleNamespace + :param k_index: The k index + :type k_index: cutlass.Int32 + :param k_tile_count: The k tile count + :type k_tile_count: cutlass.Int32 + :param load_q_producer_state: The load q producer state + :type load_q_producer_state: pipeline.PipelineState + :param load_k_producer_state: The load k producer state + :type load_k_producer_state: pipeline.PipelineState + + :return: The load q producer state and load k producer state + :rtype: tuple[pipeline.PipelineState, pipeline.PipelineState] + """ + # page table + mPT = common_params.mPT[None, common_params.blk_coord[2]] + + # Flatten divide and partition global tensors for QK TMA load + # (bM, bK, rM, rK, rL) + mma_qk_tiler_mk = cute.select(self.mma_qk_tiler, mode=[0, 2]) + gQL = cute.flat_divide(qk_params.mQL, mma_qk_tiler_mk) + mma_qk_tiler_mk_rope = cute.select(self.mma_qk_rope_tiler, mode=[0, 2]) + gQR = cute.flat_divide(qk_params.mQR, mma_qk_tiler_mk_rope) + + thr_mma_qk = qk_params.tiled_mma_qk.get_slice( + common_params.blk_coord[0] % cute.size(qk_params.tiled_mma_qk.thr_id) + ) + tSgQL = thr_mma_qk.partition_A(gQL) + tSgQR = thr_mma_qk.partition_A(gQR) + + cta_m = min( + qk_params.tiled_mma_qk.op.shape_mnk[0] + // qk_params.tiled_mma_qk.thr_id.shape, + self.page_size, + ) + page_tile_size = min(self.page_size, cta_m) + gCL = cute.tiled_divide(qk_params.mCL, (page_tile_size, self.mma_qk_tiler[2])) + tSgCL = ( + gCL[ + None, + common_params.blk_coord[0] % qk_params.tiled_mma_qk.thr_id.shape, + None, + None, + ] + if cta_m < self.page_size + else gCL[None, 0, None, None] + ) + gKR = cute.tiled_divide( + qk_params.mKR, (page_tile_size, self.mma_qk_rope_tiler[2]) + ) + tSgKR = ( + gKR[ + None, + common_params.blk_coord[0] % qk_params.tiled_mma_qk.thr_id.shape, + None, + None, + ] + if cta_m < self.page_size + else gKR[None, 0, None, None] + ) + # tma partition for q, k latent/rope + + # smem: ((atom_v, rest_v), STAGE) + # gmem: ((atom_v, rest_v), RestM, RestK, RestL) + tQsQ, tQLgQL_mkl = cpasync.tma_partition( + qk_params.tma_atom_q_latent, + 0, + cute.make_layout(1), + cute.group_modes(qk_params.sQ, 0, 3), + cute.group_modes(tSgQL, 0, 3), + ) + + tQsQ_rope, tQRgQR_mkl = cpasync.tma_partition( + qk_params.tma_atom_q_rope, + 0, + cute.make_layout(1), + cute.group_modes(qk_params.sQ_rope, 0, 3), + cute.group_modes(tSgQR, 0, 3), + ) + tKCsKC, tCLgCL = cpasync.tma_partition( + qk_params.tma_atom_c_latent, + 0, + cute.make_layout(1), + qk_params.sKC, + tSgCL, + ) + + tKCsKC_rope, tKRgKR = cpasync.tma_partition( + qk_params.tma_atom_c_rope, + 0, + cute.make_layout(1), + qk_params.sKC_rope, + tSgKR, + ) + + tQLgQL = tQLgQL_mkl[ + None, None, None, common_params.blk_coord[1], common_params.blk_coord[2] + ] + tQRgQR = tQRgQR_mkl[ + None, None, None, common_params.blk_coord[1], common_params.blk_coord[2] + ] + + # set extra params + common_params.mPT = mPT + qk_params.tQLgQL = tQLgQL + qk_params.tQRgQR = tQRgQR + qk_params.tCLgCL = tCLgCL + qk_params.tKRgKR = tKRgKR + qk_params.tQsQ = tQsQ + qk_params.tQsQ_rope = tQsQ_rope + qk_params.tKCsKC = tKCsKC + qk_params.tKCsKC_rope = tKCsKC_rope + + k_tile_count_init = k_tile_count + while k_tile_count > 0: + load_q_producer_state, load_k_producer_state = self.load_tma_qk_one_k_tile( + common_params, + qk_params, + k_index, + k_tile_count, + load_q_producer_state, + load_k_producer_state, + load_q=k_tile_count_init == k_tile_count, + ) + k_index += 1 + k_tile_count -= 1 + + return load_q_producer_state, load_k_producer_state + + @cute.jit + def load_tma_v( + self, + common_params: SimpleNamespace, + v_params: SimpleNamespace, + k_index: cutlass.Int32, + k_tile_count: cutlass.Int32, + load_v_producer_state: pipeline.PipelineState, + ) -> pipeline.PipelineState: + """Load wrap to load V tensors. Updates the load v producer state. + + :param common_params: The common parameters + :type common_params: SimpleNamespace + :param v_params: The v parameters + :type v_params: SimpleNamespace + :param k_index: The k index + :type k_index: cutlass.Int32 + :param k_tile_count: The k tile count + :type k_tile_count: cutlass.Int32 + :param load_v_producer_state: The load v producer state + :type load_v_producer_state: pipeline.PipelineState + + :return: The load v producer state + :rtype: pipeline.PipelineState + """ + # page table + mPT = common_params.mPT[None, common_params.blk_coord[2]] + + # Flatten divide and partition global tensors for V TMA load + page_tile_size = min(self.page_size, self.mma_pv_tiler[2]) + gCLT = cute.flat_divide(v_params.mCLT, (self.mma_pv_tiler[1], page_tile_size)) + cta_n = self.mma_pv_tiler[1] // v_params.tiled_mma_pv.thr_id.shape + gCLT = cute.logical_divide(gCLT, (cta_n,))[ + (None, common_params.blk_coord[0]), None, None, None, None + ] + tOgCLT = cute.tiled_divide(gCLT, (cta_n, page_tile_size)) + tOgCLT = tOgCLT[None, 0, 0, None, None, None] + # tma partition for vc + # smem: ((atom_v, rest_v), STAGE) + # gmem: ((atom_v, rest_v), RestM, RestK, RestL) + tVCsVC, tCLTgCLT = cpasync.tma_partition( + v_params.tma_atom_c_latent_transpose, + 0, + cute.make_layout(1), + v_params.sVC, + tOgCLT, + ) + + # set extra params + common_params.mPT = mPT + v_params.tCLTgCLT = tCLTgCLT + v_params.tVCsVC = tVCsVC + + while k_tile_count > 0: + load_v_producer_state = self.load_tma_v_one_k_tile( + common_params, + v_params, + k_index, + load_v_producer_state, + ) + k_index += 1 + k_tile_count -= 1 + return load_v_producer_state + + @cute.jit + def load_tma_qk_one_k_tile( + self, + common_params: SimpleNamespace, + qk_params: SimpleNamespace, + k_index: cutlass.Int32, + k_tile_count: cutlass.Int32, + load_q_producer_state: pipeline.PipelineState, + load_k_producer_state: pipeline.PipelineState, + load_q: bool, + ) -> tuple[pipeline.PipelineState, pipeline.PipelineState]: + """Load one k-tile of Q/C latent/rope tensors. Updates the load qkv producer state. + + :param common_params: The common parameters + :type common_params: SimpleNamespace + :param qk_params: The qk parameters + :type qk_params: SimpleNamespace + :param k_index: The k index + :type k_index: cutlass.Int32 + :param k_tile_count: The k tile count + :type k_tile_count: cutlass.Int32 + :param load_q_producer_state: The load q producer state + :type load_q_producer_state: pipeline.PipelineState + :param load_k_producer_state: The load kv producer state + :type load_k_producer_state: pipeline.PipelineState + :param load_q: Whether to load q + :type load_q: bool + + :return: The load q producer state and load kv producer state + :rtype: tuple[pipeline.PipelineState, pipeline.PipelineState] + """ + page_per_tile = ceil_div( + self.mma_qk_tiler[1] // self.page_size, qk_params.tiled_mma_qk.thr_id.shape + ) + k_idx = cute.make_rmem_tensor(cute.make_layout(page_per_tile), cutlass.Int32) + for i in cutlass.range_constexpr(page_per_tile): + k_idx[i] = ( + common_params.mPT[k_index] + if self.mma_qk_tiler[1] // self.page_size == 1 + else common_params.mPT[ + ( + k_index * qk_params.tiled_mma_qk.thr_id.shape + + common_params.blk_coord[0] + ) + * page_per_tile + + i + ] + ) + # load q once at first iteration + load_q_pipeline = common_params.load_q_pipeline + if load_q: + # get the mbar ptr from pipeline. + tma_bar_ptr = load_q_pipeline.producer_get_barrier(load_q_producer_state) + # expect the extra bytes for q. + load_q_pipeline.producer_acquire(load_q_producer_state) + for i in cutlass.range_constexpr(self.iterations_qk_latent): + # load q latent + cute.copy( + qk_params.tma_atom_q_latent, + qk_params.tQLgQL[None, 0, i], + qk_params.tQsQ[None, (i, 0)], + tma_bar_ptr=tma_bar_ptr, + ) + for i in cutlass.range_constexpr(self.iterations_qk_rope): + # load q rope + cute.copy( + qk_params.tma_atom_q_rope, + qk_params.tQRgQR[None, 0, i], + qk_params.tQsQ_rope[None, i], + tma_bar_ptr=tma_bar_ptr, + ) + load_q_producer_state.advance() + # get the mbar ptr from pipeline. + tma_bar_ptr = common_params.load_k_pipeline.producer_get_barrier( + load_k_producer_state + ) + common_params.load_k_pipeline.producer_acquire(load_k_producer_state) + for i in range(self.iterations_qk_latent): + for k in range(page_per_tile): + # load k latent + cute.copy( + qk_params.tma_atom_c_latent, + qk_params.tCLgCL[None, i, k_idx[k]], + qk_params.tKCsKC[None, k, 0, (i, load_k_producer_state.index)], + tma_bar_ptr=tma_bar_ptr, + ) + + for i in cutlass.range_constexpr(self.iterations_qk_rope): + for k in cutlass.range_constexpr(page_per_tile): + # load k rope + cute.copy( + qk_params.tma_atom_c_rope, + qk_params.tKRgKR[None, i, k_idx[k]], + qk_params.tKCsKC_rope[None, k, 0, load_k_producer_state.index], + tma_bar_ptr=tma_bar_ptr, + ) + load_k_producer_state.advance() + + return load_q_producer_state, load_k_producer_state + + @cute.jit + def load_tma_v_one_k_tile( + self, + common_params: SimpleNamespace, + v_params: SimpleNamespace, + k_index: cutlass.Int32, + load_v_producer_state: pipeline.PipelineState, + ) -> pipeline.PipelineState: + """Load one k-tile of compressed latent transpose tensor(v). Updates the load qkv producer state. + + :param common_params: The common parameters + :type common_params: SimpleNamespace + :param v_params: The load tma v parameters + :type v_params: SimpleNamespace + :param k_index: The k index + :type k_index: cutlass.Int32 + :param load_v_producer_state: The load v producer state + :type load_v_producer_state: pipeline.PipelineState + + :return: The load qkv producer state + :rtype: pipeline.PipelineState + """ + page_per_tile = self.mma_pv_tiler[2] * self.iterations_pv_k // self.page_size + page_per_subtile = ceil_div(page_per_tile, self.iterations_pv_k) + k_idx = cute.make_rmem_tensor(cute.make_layout(page_per_tile), cutlass.Int32) + for i in cutlass.range_constexpr(page_per_tile): + k_idx[i] = ( + common_params.mPT[k_index] + if page_per_tile == 1 + else common_params.mPT[k_index * page_per_tile + i] + ) + # get the mbar ptr from pipeline. + tma_bar_ptr = common_params.load_v_pipeline.producer_get_barrier( + load_v_producer_state + ) + common_params.load_v_pipeline.producer_acquire(load_v_producer_state) + for j in cutlass.range_constexpr(self.iterations_pv_n): + for i in cutlass.range_constexpr(self.iterations_pv_k): + if cutlass.const_expr(page_per_tile > 1): + for k in cutlass.range_constexpr(page_per_subtile): + k_idx_i = k_idx[k + i * page_per_subtile] + cute.copy( + v_params.tma_atom_c_latent_transpose, + v_params.tCLTgCLT[None, j, 0, k_idx_i], + v_params.tVCsVC[ + None, 0, k, ((j, i), load_v_producer_state.index) + ], + tma_bar_ptr=tma_bar_ptr, + ) + else: + cute.copy( + v_params.tma_atom_c_latent_transpose, + v_params.tCLTgCLT[None, j, i, k_idx[0]], + v_params.tVCsVC[ + None, 0, 0, ((j, i), load_v_producer_state.index) + ], + tma_bar_ptr=tma_bar_ptr, + ) + load_v_producer_state.advance() + return load_v_producer_state + + @cute.jit + def mma( + self, + common_params: SimpleNamespace, + qk_params: SimpleNamespace, + pv_params: SimpleNamespace, + k_tile_count: cutlass.Int32, + tiled_mma_qk: cute.TiledMma, + tiled_mma_pv: cute.TiledMma, + load_q_consumer_state: pipeline.PipelineState, + load_k_consumer_state: pipeline.PipelineState, + load_v_consumer_state: pipeline.PipelineState, + mma_s_producer_state: pipeline.PipelineState, + p_mma_consumer_state: pipeline.PipelineState, + mma_o_producer_state: pipeline.PipelineState, + ) -> tuple[ + cute.TiledMma, + cute.TiledMma, + pipeline.PipelineState, + pipeline.PipelineState, + pipeline.PipelineState, + pipeline.PipelineState, + pipeline.PipelineState, + ]: + """MMA warp to compute the result of Q*K^T and P*V. Updates the tiled mma and pipeline states. + + :param common_params: The common parameters for mma qk and pv + :type common_params: SimpleNamespace + :param qk_params: The mma qk parameters + :type qk_params: SimpleNamespace + :param pv_params: The mma pv parameters + :type pv_params: SimpleNamespace + :param k_tile_count: The k tile count + :type k_tile_count: cutlass.Int32 + :param tiled_mma_qk: The tiled mma qk + :type tiled_mma_qk: cute.TiledMma + :param tiled_mma_pv: The tiled mma pv + :type tiled_mma_pv: cute.TiledMma + :param load_q_consumer_state: The load q consumer state + :type load_q_consumer_state: pipeline.PipelineState + :param load_k_consumer_state: The load k consumer state + :type load_k_consumer_state: pipeline.PipelineState + :param load_v_consumer_state: The load v consumer state + :type load_v_consumer_state: pipeline.PipelineState + :param mma_s_producer_state: The mma s producer state + :type mma_s_producer_state: pipeline.PipelineState + :param p_mma_consumer_state: The p mma consumer state + :type p_mma_consumer_state: pipeline.PipelineState + :param mma_o_producer_state: The mma o producer state + :type mma_o_producer_state: pipeline.PipelineState + + :return: The tiled mma qk, the tiled mma pv, the load q consumer state, the load k consumer state, the load v consumer state, the mma s producer state, the p mma consumer state, and the mma o producer state + :rtype: tuple[cute.TiledMma, cute.TiledMma, pipeline.PipelineState, pipeline.PipelineState, pipeline.PipelineState, pipeline.PipelineState, pipeline.PipelineState, pipeline.PipelineState] + """ + + tSrQ = tiled_mma_qk.make_fragment_A(qk_params.sQ) + tSrQ_rope = tiled_mma_qk.make_fragment_A(qk_params.sQ_rope) + tSrKC = tiled_mma_qk.make_fragment_B(qk_params.sKC) + tSrKC_rope = tiled_mma_qk.make_fragment_B(qk_params.sKC_rope) + tOrP = tiled_mma_pv.make_fragment_A(pv_params.sP) + tOrVC = tiled_mma_pv.make_fragment_B(pv_params.sVC) + + tStS_shape = tiled_mma_qk.partition_shape_C( + cute.select(self.mma_qk_tiler, mode=[0, 1]) + ) + tStS_staged_fake = tiled_mma_qk.make_fragment_C( + cute.append(tStS_shape, self.mma_s_stage) + ) + # use real tmem ptr for tStS + tStS_staged = cute.make_tensor(common_params.tmem_ptr, tStS_staged_fake.layout) + tOtO_shape = tiled_mma_pv.partition_shape_C( + cute.select(self.mma_pv_tiler, mode=[0, 1]) + ) + # mma O has 1 stage. + tOtO = tiled_mma_pv.make_fragment_C(tOtO_shape) + tOtO_layout = cute.append( + tOtO.layout, + cute.make_layout( + common_params.L // self.mma_pv_tiler[1], + stride=self.mma_pv_tiler[1] // self.warps_in_n, + ), + ) + tOtO_staged = cute.make_tensor( + tStS_staged.iterator + self.tmem_o_offset, tOtO_layout + ) + + # set more parameters + qk_params.tSrQ = tSrQ + qk_params.tSrQ_rope = tSrQ_rope + qk_params.tSrKC = tSrKC + qk_params.tSrKC_rope = tSrKC_rope + qk_params.tStS_staged = tStS_staged + pv_params.tOrP = tOrP + pv_params.tOrVC = tOrVC + pv_params.tOtO_staged = tOtO_staged + + # mma O accumulates on K, so the accumlate flag is set to False once before all K blocks. + tiled_mma_pv.set(tcgen05.Field.ACCUMULATE, False) + load_q_pipeline = common_params.load_q_pipeline + if common_params.is_leader_cta: + load_q_release_state = load_q_consumer_state.clone() + ( + tiled_mma_qk, + load_q_consumer_state, + load_k_consumer_state, + mma_s_producer_state, + ) = self.mma_qk( + common_params, + qk_params, + tiled_mma_qk, + load_q_consumer_state, + load_k_consumer_state, + mma_s_producer_state, + wait_q=True, + ) + k_tile_count -= 1 + + while k_tile_count > 0: + ( + tiled_mma_qk, + load_q_consumer_state, + load_k_consumer_state, + mma_s_producer_state, + ) = self.mma_qk( + common_params, + qk_params, + tiled_mma_qk, + load_q_consumer_state, + load_k_consumer_state, + mma_s_producer_state, + wait_q=False, + ) + ( + tiled_mma_pv, + load_v_consumer_state, + p_mma_consumer_state, + mma_o_producer_state, + ) = self.mma_pv( + common_params, + pv_params, + tiled_mma_pv, + load_v_consumer_state, + p_mma_consumer_state, + mma_o_producer_state, + ) + k_tile_count -= 1 + # release q consumer states + load_q_pipeline.consumer_release(load_q_release_state) + load_q_release_state.advance() + ( + tiled_mma_pv, + load_v_consumer_state, + p_mma_consumer_state, + mma_o_producer_state, + ) = self.mma_pv( + common_params, + pv_params, + tiled_mma_pv, + load_v_consumer_state, + p_mma_consumer_state, + mma_o_producer_state, + ) + + return ( + tiled_mma_qk, + tiled_mma_pv, + load_q_consumer_state, + load_k_consumer_state, + load_v_consumer_state, + mma_s_producer_state, + p_mma_consumer_state, + mma_o_producer_state, + ) + + @cute.jit + def mma_qk( + self, + common_params: SimpleNamespace, + qk_params: SimpleNamespace, + tiled_mma_qk: cute.TiledMma, + load_q_consumer_state: pipeline.PipelineState, + load_k_consumer_state: pipeline.PipelineState, + mma_s_producer_state: pipeline.PipelineState, + wait_q: bool, + ) -> tuple[ + cute.TiledMma, + pipeline.PipelineState, + pipeline.PipelineState, + pipeline.PipelineState, + ]: + """Compute one k-tile of mma for Q*K^T. Updates the tiled MMA QK and pipeline states. + + :param qk_params: The qk parameters + :type qk_params: SimpleNamespace + :param tiled_mma_qk: The tiled mma qk + :type tiled_mma_qk: cute.TiledMma + :param load_q_consumer_state: The load q consumer state + :type load_q_consumer_state: pipeline.PipelineState + :param load_k_consumer_state: The load k consumer state + :type load_k_consumer_state: pipeline.PipelineState + :param mma_s_producer_state: The mma s producer state + :type mma_s_producer_state: pipeline.PipelineState + + :return: The tiled mma qk, the load q consumer state, the load k consumer state, and the mma s producer state + :rtype: tuple[cute.TiledMma, pipeline.PipelineState, pipeline.PipelineState, pipeline.PipelineState] + """ + tStS = qk_params.tStS_staged[None, None, None, mma_s_producer_state.index] + + qk_params.mma_s_pipeline.producer_acquire(mma_s_producer_state) + tiled_mma_qk.set(tcgen05.Field.ACCUMULATE, False) + load_q_pipeline = common_params.load_q_pipeline + load_k_pipeline = common_params.load_k_pipeline + if cutlass.const_expr(wait_q): + load_q_pipeline.consumer_wait(load_q_consumer_state) + load_k_pipeline.consumer_wait(load_k_consumer_state) + for q_stage in range(self.iterations_qk_latent): + kc_stage = load_k_consumer_state.index + for k_block in cutlass.range_constexpr(cute.size(qk_params.tSrQ.shape[2])): + cute.gemm( + tiled_mma_qk, + tStS, + qk_params.tSrQ[None, None, k_block, (q_stage, 0)], + qk_params.tSrKC[None, None, k_block, (q_stage, kc_stage)], + tStS, + ) + tiled_mma_qk.set(tcgen05.Field.ACCUMULATE, True) + + for q_stage in range(self.iterations_qk_rope): + kc_stage = load_k_consumer_state.index + for k_block in cutlass.range_constexpr( + self.rope_dim // tiled_mma_qk.shape_mnk[2] + ): + cute.gemm( + tiled_mma_qk, + tStS, + qk_params.tSrQ_rope[None, None, k_block, q_stage], + qk_params.tSrKC_rope[None, None, k_block, kc_stage], + tStS, + ) + tiled_mma_qk.set(tcgen05.Field.ACCUMULATE, True) + load_k_pipeline.consumer_release(load_k_consumer_state) + load_k_consumer_state.advance() + if cutlass.const_expr(wait_q): + load_q_consumer_state.advance() + + qk_params.mma_s_pipeline.producer_commit(mma_s_producer_state) + mma_s_producer_state.advance() + return ( + tiled_mma_qk, + load_q_consumer_state, + load_k_consumer_state, + mma_s_producer_state, + ) + + @cute.jit + def mma_pv( + self, + common_params: SimpleNamespace, + pv_params: SimpleNamespace, + tiled_mma_pv: cute.TiledMma, + load_v_consumer_state: pipeline.PipelineState, + p_mma_consumer_state: pipeline.PipelineState, + mma_o_producer_state: pipeline.PipelineState, + ) -> tuple[ + cute.TiledMma, + pipeline.PipelineState, + pipeline.PipelineState, + pipeline.PipelineState, + ]: + """Compute one k-tile of mma for P*V. Updates the tiled mma pv and pipeline states. + + :param common_params: The common parameters + :type common_params: SimpleNamespace + :param pv_params: The pv parameters + :type pv_params: SimpleNamespace + :param tiled_mma_pv: The tiled mma pv + :type tiled_mma_pv: cute.TiledMma + :param load_v_consumer_state: The load v consumer state + :type load_v_consumer_state: pipeline.PipelineState + :param p_mma_consumer_state: The P MMA consumer state + :type p_mma_consumer_state: pipeline.PipelineState + :param mma_o_producer_state: The MMA o producer state + :type mma_o_producer_state: pipeline.PipelineState + + :return: The tiled mma pv, the load v consumer state, the P MMA consumer state, and the MMA o producer state + :rtype: tuple[cute.TiledMma, pipeline.PipelineState, pipeline.PipelineState, pipeline.PipelineState] + """ + + pv_params.p_mma_pipeline.consumer_wait(p_mma_consumer_state) + load_v_pipeline = common_params.load_v_pipeline + accumulate_flag = tiled_mma_pv.get(tcgen05.Field.ACCUMULATE) + mma_o_pipeline = pv_params.mma_o_pipeline + + load_v_pipeline.consumer_wait(load_v_consumer_state) + vc_stage = load_v_consumer_state.index + for acc_stage in range(self.iterations_pv_n): + mma_o_pipeline.producer_acquire(mma_o_producer_state) + tiled_mma_pv.set(tcgen05.Field.ACCUMULATE, accumulate_flag) + for p_stage in range(self.iterations_pv_k): + tOtO = pv_params.tOtO_staged[None, None, None, acc_stage] + for k_block in cutlass.range_constexpr(pv_params.tOrP.shape[2]): + cute.gemm( + tiled_mma_pv, + tOtO, + pv_params.tOrP[ + None, + None, + k_block, + (p_stage, p_mma_consumer_state.index), + ], + pv_params.tOrVC[ + None, None, k_block, ((acc_stage, p_stage), vc_stage) + ], + tOtO, + ) + tiled_mma_pv.set(tcgen05.Field.ACCUMULATE, True) + + mma_o_pipeline.producer_commit(mma_o_producer_state) + mma_o_producer_state.advance() + load_v_pipeline.consumer_release(load_v_consumer_state) + load_v_consumer_state.advance() + pv_params.p_mma_pipeline.consumer_release(p_mma_consumer_state) + p_mma_consumer_state.advance() + + return ( + tiled_mma_pv, + load_v_consumer_state, + p_mma_consumer_state, + mma_o_producer_state, + ) + + @cute.jit + def compute( + self, + common_params: SimpleNamespace, + softmax_params: SimpleNamespace, + k_index: cutlass.Int32, + k_tile_count: cutlass.Int32, + mma_s_consumer_state: pipeline.PipelineState, + p_mma_producer_state: pipeline.PipelineState, + p_cor_producer_state: pipeline.PipelineState, + ) -> tuple[pipeline.PipelineState, pipeline.PipelineState, pipeline.PipelineState]: + """Compute warp to compute the result of softmax, rescale, and epilogue. Updates the related pipeline states. + + :param common_params: The common parameters + :type common_params: SimpleNamespace + :param softmax_params: The softmax parameters + :type softmax_params: SimpleNamespace + :param k_index: The index of the k-tile + :type k_index: cutlass.Int32 + :param k_tile_count: The number of k-tiles + :type k_tile_count: cutlass.Int32 + :param mma_s_consumer_state: The MMA s consumer state + :type mma_s_consumer_state: pipeline.PipelineState + :param p_mma_producer_state: The P MMA producer state + :type p_mma_producer_state: pipeline.PipelineState + :param p_cor_producer_state: The P correction producer state + :type p_cor_producer_state: pipeline.PipelineState + + :return: The MMA s consumer state, the P MMA producer state, and the P correction producer state + :rtype: tuple[pipeline.PipelineState, pipeline.PipelineState, pipeline.PipelineState] + """ + + k_tile_total = cute.ceil_div(common_params.K, self.mma_qk_tiler[1]) + + row_max = -self.acc_dtype.inf + row_sum = self.acc_dtype(0) + correction_factor = self.acc_dtype(1) + common_params.p_cor_pipeline.producer_acquire(p_cor_producer_state) + + # no mask applied + while k_tile_count > 1: + ( + mma_s_consumer_state, + p_mma_producer_state, + p_cor_producer_state, + row_max, + row_sum, + correction_factor, + ) = self.softmax( + common_params, + softmax_params, + k_index, + mma_s_consumer_state, + p_mma_producer_state, + p_cor_producer_state, + row_max, + row_sum, + correction_factor, + False, + False, + ) + k_index = k_index + 1 + k_tile_count = k_tile_count - 1 + + # mask applied + if cutlass.const_expr(common_params.mAccO is not None): + ( + mma_s_consumer_state, + p_mma_producer_state, + p_cor_producer_state, + row_max, + row_sum, + correction_factor, + ) = self.softmax( + common_params, + softmax_params, + k_index, + mma_s_consumer_state, + p_mma_producer_state, + p_cor_producer_state, + row_max, + row_sum, + correction_factor, + k_index == k_tile_total - 1, + True, + ) + else: + ( + mma_s_consumer_state, + p_mma_producer_state, + p_cor_producer_state, + row_max, + row_sum, + correction_factor, + ) = self.softmax( + common_params, + softmax_params, + k_index, + mma_s_consumer_state, + p_mma_producer_state, + p_cor_producer_state, + row_max, + row_sum, + correction_factor, + True, + True, + ) + + return mma_s_consumer_state, p_mma_producer_state, p_cor_producer_state + + @cute.jit + def correction( + self, + common_params: SimpleNamespace, + epilogue_params: SimpleNamespace, + k_tile_count: cutlass.Int32, + p_cor_consumer_state: pipeline.PipelineState, + mma_o_consumer_state: pipeline.PipelineState, + ) -> tuple[pipeline.PipelineState, pipeline.PipelineState]: + """Compute warp to compute the result of softmax, rescale, and epilogue. Updates the related pipeline states. + + :param common_params: The common parameters + :type common_params: SimpleNamespace + :param epilogue_params: The epilogue parameters + :type epilogue_params: SimpleNamespace + :param k_index: The index of the k-tile + :type k_index: cutlass.Int32 + :param k_tile_count: The number of k-tiles + :type k_tile_count: cutlass.Int32 + :param p_cor_consumer_state: The P correction consumer state + :type p_cor_consumer_state: pipeline.PipelineState + :param mma_o_consumer_state: The MMA o consumer state + :type mma_o_consumer_state: pipeline.PipelineState + + :return: The P correction consumer state, and the MMA o consumer state + :rtype: tuple[pipeline.PipelineState, pipeline.PipelineState] + """ + + k_tile_count_init = k_tile_count + while k_tile_count > 0: + p_cor_consumer_state, row_sum, row_max, correction_factor, no_correction = ( + self.get_correction_factor(common_params, p_cor_consumer_state) + ) + if k_tile_count_init != k_tile_count: + mma_o_consumer_state = self.rescale( + common_params, + mma_o_consumer_state, + correction_factor, + no_correction, + ) + k_tile_count = k_tile_count - 1 + if k_tile_count == 0: + mma_o_consumer_state = self.epilogue( + common_params, + epilogue_params, + mma_o_consumer_state, + row_sum, + row_max, + ) + return p_cor_consumer_state, mma_o_consumer_state + + @cute.jit + def exchange_p_cor_metadata( + self, + common_params: SimpleNamespace, + softmax_params: SimpleNamespace, + correction_factor: cutlass.Float32, + row_sum: cutlass.Float32, + row_max: cutlass.Float32, + row_max_new: cutlass.Float32, + tAcc: cute.Tensor, + tidx: cutlass.Int32, + p_cor_producer_state: pipeline.PipelineState, + ) -> tuple[pipeline.PipelineState, cutlass.Float32]: + """Compute the correction factor for the last k tile.""" + no_correction = 0 + if ( + row_max_new - row_max + ) * softmax_params.softmax_scale_log2 <= self.skip_correction_threshold: + no_correction = 1 + row_max_new = row_max + + # pad for 4x32b + corr_layout = cute.make_layout( + (tAcc.shape[0], (4, tAcc.shape[1][1]), self.mma_s_stage), + stride=(tAcc.stride[0], (1, tAcc.stride[1][1]), 4), + ) + tCor = cute.make_tensor( + common_params.tmem_ptr + self.correction_factor_offset, + corr_layout, + ) + cCor = cute.make_identity_tensor(tCor.shape) + corr_tmem_store_atom = cute.make_copy_atom( + tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(4)), self.acc_dtype + ) + corr_tmem_store_tiled_copy = tcgen05.make_tmem_copy(corr_tmem_store_atom, tCor) + corr_tmem_store_thr_copy = corr_tmem_store_tiled_copy.get_slice(tidx) + cCor_for_copy = corr_tmem_store_thr_copy.partition_S(cCor) + tCor_for_copy = corr_tmem_store_thr_copy.partition_D(tCor) + rCor = cute.make_fragment_like( + cCor_for_copy[None, None, None, 0], self.acc_dtype + ) + rCor_int = cute.make_tensor( + cute.recast_ptr(rCor.iterator, dtype=cutlass.Int32), rCor.layout + ) + rCor[0] = row_sum + rCor[1] = row_max_new + rCor[2] = correction_factor + rCor_int[3] = no_correction + + cute.copy( + corr_tmem_store_tiled_copy, + rCor, + tCor_for_copy[None, None, None, p_cor_producer_state.index], + ) + # fence between tmem store and correction warp + cute.arch.fence_view_async_tmem_store() + common_params.p_cor_pipeline.producer_commit(p_cor_producer_state) + p_cor_producer_state.advance() + return p_cor_producer_state, row_max_new + + @cute.jit + def softmax( + self, + common_params: SimpleNamespace, + softmax_params: SimpleNamespace, + k_index: cutlass.Int32, + mma_s_consumer_state: pipeline.PipelineState, + p_mma_producer_state: pipeline.PipelineState, + p_cor_producer_state: pipeline.PipelineState, + row_max: cutlass.Float32, + row_sum: cutlass.Float32, + correction_factor: cutlass.Float32, + is_last_tile: bool, + is_local_last_tile: cutlass.Boolean, + ) -> tuple[ + pipeline.PipelineState, + pipeline.PipelineState, + pipeline.PipelineState, + cutlass.Float32, + cutlass.Float32, + cutlass.Float32, + ]: + """Softmax for one k-tile. Updates the related pipeline states and returns the computed results. + + :param common_params: The common parameters + :type common_params: SimpleNamespace + :param softmax_params: The softmax parameters + :type softmax_params: SimpleNamespace + :param k_index: The index of the k-tile + :type k_index: cutlass.Int32 + :param mma_s_consumer_state: The MMA s consumer state + :type mma_s_consumer_state: pipeline.PipelineState + :param p_mma_producer_state: The P MMA producer state + :type p_mma_producer_state: pipeline.PipelineState + :param p_cor_producer_state: The P correction producer state + :type p_cor_producer_state: pipeline.PipelineState + :param row_max: The row max + :type row_max: cutlass.Float32 + :param row_sum: The row sum + :type row_sum: cutlass.Float32 + :param correction_factor: The correction factor + :type correction_factor: cutlass.Float32 + :param is_last_tile: Whether the last tile + :type is_last_tile: bool + :param is_local_last_tile: Whether the last tile is local + :type is_local_last_tile: cutlass.Boolean + + :return: The MMA s consumer state, the P MMA producer state, the P correction producer state, the row max, the row sum, and the correction factor + :rtype: tuple[pipeline.PipelineState, pipeline.PipelineState, pipeline.PipelineState, cutlass.Float32, cutlass.Float32, cutlass.Float32] + """ + + softmax_params.p_mma_pipeline.producer_acquire(p_mma_producer_state) + softmax_params.mma_s_pipeline.consumer_wait(mma_s_consumer_state) + + # load S from tmem + tStS_shape = softmax_params.tiled_mma_qk.partition_shape_C( + cute.select(self.mma_qk_tiler, mode=[0, 1]) + ) + tStS_staged_fake = softmax_params.tiled_mma_qk.make_fragment_C( + cute.append(tStS_shape, self.mma_s_stage) + ) + tStS_staged = cute.make_tensor(common_params.tmem_ptr, tStS_staged_fake.layout) + tStS = tStS_staged[None, None, None, mma_s_consumer_state.index] + + tAcc = tStS[(None, None), 0, 0] + cta_qk_tiler = ( + self.mma_qk_tiler[0] // self.cluster_shape_mnk[0], + self.mma_qk_tiler[1], + self.mma_qk_tiler[2], + ) + cS = cute.make_identity_tensor(cute.select(cta_qk_tiler, mode=[0, 1])) + + tmem_load_atom = cute.make_copy_atom( + tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(32)), self.acc_dtype + ) + tmem_tiled_copy = tcgen05.make_tmem_copy(tmem_load_atom, tAcc) + + tidx = common_params.tidx % (self.num_compute_warps * self.threads_per_warp) + + tmem_thr_copy = tmem_tiled_copy.get_slice(tidx) + tTR_tAcc = tmem_thr_copy.partition_S(tAcc) + tTR_tS = tmem_thr_copy.partition_D(cS) + + tTR_rAcc = cute.make_fragment_like(tTR_tS, self.acc_dtype) + + row_max_new = row_max + arch = BaseDSL._get_dsl().get_arch_enum() + if cutlass.const_expr(arch >= Arch.sm_100 and arch <= Arch.sm_100f): + cute.copy(tmem_tiled_copy, tTR_tAcc, tTR_rAcc) + for i in cutlass.range_constexpr(cute.size(tTR_rAcc)): + if is_last_tile: + tTR_rAcc[i] = ( + tTR_rAcc[i] + if cute.elem_less( + tTR_tS[i][1] + self.mma_qk_tiler[1] * k_index, + common_params.K, + ) + else -self.acc_dtype.inf + ) + # reduction for row_max + row_max_new = tTR_rAcc.load().reduce(cute.ReductionOp.MAX, row_max_new, 0) + elif cutlass.const_expr(arch >= Arch.sm_103 and arch <= Arch.sm_103f): + tmem_load_red_atom = cute.make_copy_atom( + tcgen05.copy.LdRed32x32bOp( + tcgen05.copy.Repetition(64), redOp=tcgen05.TmemLoadRedOp.MAX + ), + self.acc_dtype, + ) + tmem_red_tiled_copy = tcgen05.make_tmem_copy(tmem_load_red_atom, tAcc) + tmem_red_thr_copy = tmem_red_tiled_copy.get_slice(tidx) + tTR_tAcc_red = tmem_red_thr_copy.partition_S(tAcc) + tTR_tS_red = tmem_red_thr_copy.partition_D(cS) + tTR_rAcc_red = cute.make_fragment_like(tTR_tS_red, self.acc_dtype) + tTR_rMax = cute.make_rmem_tensor( + cute.make_layout((1, tTR_tS_red.shape[1], tTR_tS_red.shape[2])), + self.acc_dtype, + ) + cute.copy( + tmem_red_tiled_copy, + tTR_tAcc_red, + (tTR_rAcc_red, tTR_rMax), + ) + tTR_rAcc = cute.make_tensor(tTR_rAcc_red.iterator, tTR_rAcc.layout) + if is_last_tile: + for i in cutlass.range_constexpr(cute.size(tTR_rAcc)): + tTR_rAcc[i] = ( + tTR_rAcc[i] + if cute.elem_less( + tTR_tS[i][1] + self.mma_qk_tiler[1] * k_index, + common_params.K, + ) + else -self.acc_dtype.inf + ) + # reduction for row_max + row_max_new = tTR_rAcc.load().reduce( + cute.ReductionOp.MAX, row_max_new, 0 + ) + else: + row_max_new = cute.arch.fmax(row_max_new, tTR_rMax[0]) + + # if warps in N is 2, reduce row_max across warps (0, 1) and (2, 3) + if cutlass.const_expr(self.warps_in_n == 2): + common_params.smem_exchange[tidx] = row_max_new + self.softmax_exchange_sync_bar.wait() + row_max_new = cute.arch.fmax( + row_max_new, + common_params.smem_exchange[ + (tidx + 64) % (self.num_compute_warps * self.threads_per_warp) + ], + ) + + # find correction factor + correction_factor = cute.math.exp2( + (row_max - row_max_new) * softmax_params.softmax_scale_log2, fastmath=True + ) + # split kv case + if cutlass.const_expr(not is_local_last_tile): + p_cor_producer_state, row_max_new = self.exchange_p_cor_metadata( + common_params, + softmax_params, + correction_factor, + row_sum, + row_max, + row_max_new, + tAcc, + tidx, + p_cor_producer_state, + ) + + # softmax + fma_b = softmax_params.softmax_scale_log2 + fma_c = (0.0 - row_max_new) * softmax_params.softmax_scale_log2 + + for i in cutlass.range(cute.size(tTR_rAcc), vectorize=True, unroll_full=True): + tTR_rAcc[i] = tTR_rAcc[i] * fma_b + fma_c + tTR_rAcc[i] = cute.math.exp2(tTR_rAcc[i], fastmath=True) + + tTR_rS = cute.make_fragment_like(tTR_tS, self.q_dtype) + + # quantize + tTR_rS.store(tTR_rAcc.load().to(self.q_dtype)) + + # create sP + sP = softmax_params.sP[None, None, None, (None, p_mma_producer_state.index)] + sP_mk_view = cute.make_tensor( + sP.iterator, + cute.make_layout( + ( + (sP.shape[0][0], sP.shape[1]), + (sP.shape[0][1], sP.shape[2], sP.shape[3]), + ), + stride=( + (sP.stride[0][0], sP.stride[1]), + (sP.stride[0][1], sP.stride[2], sP.stride[3]), + ), + ), + ) + # change to PISL + sP_wo_swizzle_iter = cute.recast_ptr(sP.iterator, swizzle_=None) + swizzle_bits = ( + int(math.log2(self.mma_pv_tiler[2] * self.q_dtype.width // 8 // 32)) + 1 + ) + swizzle_base = 3 if self.q_dtype.width == 16 else 4 + sP_swizzle = cute.make_swizzle(swizzle_bits, swizzle_base, 3) + sP_mk_view = cute.make_tensor( + sP_wo_swizzle_iter, + cute.make_composed_layout(sP_swizzle, 0, sP_mk_view.layout), + ) + universal_copy_bits = 128 + smem_copy_atom = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), + self.q_dtype, + num_bits_per_copy=universal_copy_bits, + ) + smem_tiled_copy = cute.make_tiled_copy_D(smem_copy_atom, tmem_tiled_copy) + smem_thr_copy = smem_tiled_copy.get_slice(tidx) + rP_copy_view = smem_thr_copy.retile(tTR_rS) + sP_copy_view = smem_thr_copy.partition_D(sP_mk_view) + cute.copy(smem_tiled_copy, rP_copy_view, sP_copy_view) + + # fence between smem store and mma o + cute.arch.fence_view_async_shared() + softmax_params.p_mma_pipeline.producer_commit(p_mma_producer_state) + p_mma_producer_state.advance() + + # row_sum, using `add_packed_f32x2` to reduce the number of instructions + row_sum = row_sum * correction_factor + row_sum_vec = (0.0, 0.0) + for i in cutlass.range_constexpr(0, cute.size(tTR_rAcc), 2): + row_sum_vec = cute.arch.add_packed_f32x2( + row_sum_vec, (tTR_rAcc[i], tTR_rAcc[i + 1]) + ) + row_sum = row_sum_vec[0] + row_sum_vec[1] + row_sum + + # split kv case + if cutlass.const_expr(is_local_last_tile): + p_cor_producer_state, row_max_new = self.exchange_p_cor_metadata( + common_params, + softmax_params, + correction_factor, + row_sum, + row_max, + row_max_new, + tAcc, + tidx, + p_cor_producer_state, + ) + + # store correction factor/row_sum/row_max to tmem for correction warp + common_params.p_cor_pipeline.producer_acquire(p_cor_producer_state) + + # fence between tmem load and mma s + cute.arch.fence_view_async_tmem_load() + + softmax_params.mma_s_pipeline.consumer_release(mma_s_consumer_state) + mma_s_consumer_state.advance() + + return ( + mma_s_consumer_state, + p_mma_producer_state, + p_cor_producer_state, + row_max_new, + row_sum, + correction_factor, + ) + + @cute.jit + def _tmem_load_partition( + self, common_params: SimpleNamespace, tiled_mma_pv: cute.TiledMma, iter_n: int + ) -> tuple[ + cute.TiledMma, cute.TiledMma, cute.TiledMma, cute.TiledMma, cute.TiledMma + ]: + """Tensor memory load partition for rescale and epilogue. + + :param common_params: The common parameters + :type common_params: SimpleNamespace + :param tiled_mma_pv: The tiled mma pv + :type tiled_mma_pv: cute.TiledMma + :param iter_n: The iteration number + :type iter_n: int + + :return: The tiled mma pv, the tiled mma pv, the tiled mma pv, the tiled mma pv, the tiled mma pv + :rtype: tuple[cute.TiledMma, cute.TiledMma, cute.TiledMma, cute.TiledMma, cute.TiledMma] + """ + + tOtO_shape = tiled_mma_pv.partition_shape_C( + cute.select(self.mma_pv_tiler, mode=[0, 1]) + ) + tOtO = tiled_mma_pv.make_fragment_C(tOtO_shape) + tOtO_layout = cute.append( + tOtO.layout, + cute.make_layout( + common_params.L // self.mma_pv_tiler[1], + stride=self.mma_pv_tiler[1] // self.warps_in_n, + ), + ) + tOtO = cute.make_tensor( + common_params.tmem_ptr + self.tmem_o_offset, tOtO_layout + ) + tOtO = tOtO[None, None, None, iter_n] + + tAcc = tOtO[(None, None), 0, 0] + + tmem_load_atom = cute.make_copy_atom( + tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(32)), self.acc_dtype + ) + tmem_load_tiled_copy = tcgen05.make_tmem_copy(tmem_load_atom, tAcc) + tmem_load_thr_copy = tmem_load_tiled_copy.get_slice( + common_params.tidx % (self.num_compute_warps * self.threads_per_warp) + ) + + cta_pv_tiler = ( + self.mma_pv_tiler[0] // self.cluster_shape_mnk[0], + self.mma_pv_tiler[1], + self.mma_pv_tiler[2], + ) + # Flatten divide and partition global tensors for O + cta_pv_tiler_mn = cute.select(cta_pv_tiler, mode=[0, 1]) + + gO = None + if cutlass.const_expr(common_params.mAccO is not None): + gO = cute.local_tile( + common_params.mAccO[None, common_params.blk_coord[3], None, None, None], + cta_pv_tiler_mn, + ( + common_params.blk_coord[0], + iter_n, + common_params.blk_coord[1], + common_params.blk_coord[2], + ), + ) + cO = cute.local_tile( + cute.make_identity_tensor( + common_params.mAccO[ + None, common_params.blk_coord[3], None, None, None + ].shape + ), + cta_pv_tiler_mn, + ( + common_params.blk_coord[0], + iter_n, + common_params.blk_coord[1], + common_params.blk_coord[2], + ), + ) + else: + gO = cute.local_tile( + common_params.mO, + cta_pv_tiler_mn, + ( + common_params.blk_coord[0], + iter_n, + common_params.blk_coord[1], + common_params.blk_coord[2], + ), + ) + cO = cute.local_tile( + cute.make_identity_tensor(common_params.mO.shape), + cta_pv_tiler_mn, + ( + common_params.blk_coord[0], + iter_n, + common_params.blk_coord[1], + common_params.blk_coord[2], + ), + ) + tTR_tAcc = tmem_load_thr_copy.partition_S(tAcc) + tTR_gO = tmem_load_thr_copy.partition_D(gO) + tTR_cO = tmem_load_thr_copy.partition_D(cO) + tTR_rAcc = cute.make_fragment_like(tTR_gO, self.acc_dtype) + return tmem_load_tiled_copy, tAcc, tTR_tAcc, tTR_gO, tTR_cO, tTR_rAcc + + def get_correction_factor( + self, + common_params: SimpleNamespace, + p_cor_consumer_state: pipeline.PipelineState, + ) -> tuple[ + pipeline.PipelineState, + cutlass.Float32, + cutlass.Float32, + cutlass.Float32, + cutlass.Int32, + ]: + """Get the correction factor from the P correction consumer state. + + :param common_params: The common parameters + :type common_params: SimpleNamespace + :param p_cor_consumer_state: The P correction consumer state + :type p_cor_consumer_state: pipeline.PipelineState + + :return: The P correction consumer state, the row_sum, the row_max, and the correction factor + :rtype: tuple[pipeline.PipelineState, cutlass.Float32, cutlass.Float32, cutlass.Float32, cutlass.Int32] + """ + common_params.p_cor_pipeline.consumer_wait(p_cor_consumer_state) + tidx = common_params.tidx % (self.num_compute_warps * self.threads_per_warp) + # load correction factor + _, tAcc, _, _, _, _ = self._tmem_load_partition( + common_params, common_params.tiled_mma_pv, 0 + ) + corr_layout = cute.make_layout( + (tAcc.shape[0], (4, tAcc.shape[1][1]), self.p_cor_stage), + stride=(tAcc.stride[0], (1, tAcc.stride[1][1]), 4), + ) + tCor = cute.make_tensor( + common_params.tmem_ptr + self.correction_factor_offset, corr_layout + ) + cCor = cute.make_identity_tensor(tCor.shape) + corr_tmem_load_atom = cute.make_copy_atom( + tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(4)), self.acc_dtype + ) + corr_tmem_load_tiled_copy = tcgen05.make_tmem_copy(corr_tmem_load_atom, tCor) + corr_tmem_load_thr_copy = corr_tmem_load_tiled_copy.get_slice(tidx) + tCor_for_copy = corr_tmem_load_thr_copy.partition_S(tCor) + cCor_for_copy = corr_tmem_load_thr_copy.partition_D(cCor) + rCor = cute.make_fragment_like( + cCor_for_copy[None, None, None, 0], self.acc_dtype + ) + rCor_int = cute.make_tensor( + cute.recast_ptr(rCor.iterator, dtype=cutlass.Int32), rCor.layout + ) + cute.copy( + corr_tmem_load_tiled_copy, + tCor_for_copy[None, None, None, p_cor_consumer_state.index], + rCor, + ) + row_sum = rCor[0] + row_max = rCor[1] + correction_factor = rCor[2] + no_correction = rCor_int[3] + + common_params.p_cor_pipeline.consumer_release(p_cor_consumer_state) + p_cor_consumer_state.advance() + return p_cor_consumer_state, row_sum, row_max, correction_factor, no_correction + + @cute.jit + def rescale( + self, + common_params: SimpleNamespace, + mma_o_consumer_state: pipeline.PipelineState, + correction_factor: cutlass.Float32, + no_correction: cutlass.Int32, + ) -> pipeline.PipelineState: + """Rescale for one k-tile. Updates the related pipeline state. + + :param common_params: The common parameters + :type common_params: SimpleNamespace + :param mma_o_consumer_state: The mma o consumer state + :type mma_o_consumer_state: pipeline.PipelineState + :param correction_factor: The correction factor + :type correction_factor: cutlass.Float32 + :param no_correction: Whether to apply correction factor + :type no_correction: cutlass.Int32 + + :return: The MMA o consumer state + :rtype: pipeline.PipelineState + """ + skip_correction = cute.arch.vote_all_sync(no_correction == 1) + for iter_n in cutlass.range_constexpr(self.iterations_pv_n): + common_params.mma_o_pipeline.consumer_wait(mma_o_consumer_state) + if not skip_correction: + # tmem load tiled copy and partition results. + tmem_load_tiled_copy, tAcc, tTR_tAcc, tTR_gO, tTR_cO, tTR_rAcc = ( + self._tmem_load_partition( + common_params, common_params.tiled_mma_pv, iter_n + ) + ) + + # tmem store tiled copy + tmem_store_atom = cute.make_copy_atom( + tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(32)), self.acc_dtype + ) + tmem_store_tiled_copy = tcgen05.make_tmem_copy(tmem_store_atom, tAcc) + + # load o + cute.copy(tmem_load_tiled_copy, tTR_tAcc, tTR_rAcc) + # rescale, using `mul_packed_f32x2` to reduce the number of instructions + for i in cutlass.range( + cute.size(tTR_rAcc), vectorize=True, unroll_full=True + ): + tTR_rAcc[i] = tTR_rAcc[i] * correction_factor + + # store o to tensor memory for next k tile + cute.copy(tmem_store_tiled_copy, tTR_rAcc, tTR_tAcc) + + cute.arch.fence_view_async_tmem_store() + common_params.mma_o_pipeline.consumer_release(mma_o_consumer_state) + mma_o_consumer_state.advance() + + return mma_o_consumer_state + + @cute.jit + def epilogue( + self, + common_params: SimpleNamespace, + epilogue_params: SimpleNamespace, + mma_o_consumer_state: pipeline.PipelineState, + row_sum: cutlass.Float32, + row_max: cutlass.Float32, + ) -> pipeline.PipelineState: + """Epilogue for one k-tile. Updates the related pipeline state. + + :param common_params: The common parameters + :type common_params: SimpleNamespace + :param epilogue_params: The epilogue parameters + :type epilogue_params: SimpleNamespace + :param mma_o_consumer_state: The mma o consumer state + :type mma_o_consumer_state: pipeline.PipelineState + :param row_sum: The row sum + :type row_sum: cutlass.Float32 + :param row_max: The row max + :type row_max: cutlass.Float32 + + :return: The MMA o consumer state + :rtype: pipeline.PipelineState + """ + + tidx = common_params.tidx % (self.num_compute_warps * self.threads_per_warp) + + # exchange row_sum between warps (0, 1) and (2, 3) + if cutlass.const_expr(self.warps_in_n == 2): + common_params.smem_exchange[tidx] = row_sum + self.epilogue_exchange_sync_bar.wait() + # (64, 2) + row_sum = ( + row_sum + + common_params.smem_exchange[ + (tidx + 64) % (self.num_compute_warps * self.threads_per_warp) + ] + ) + # mma_o pipeline consumer wait + for iter_n in cutlass.range_constexpr(self.iterations_pv_n): + common_params.mma_o_pipeline.consumer_wait(mma_o_consumer_state) + # tmem load tiled copy and partition results. + tmem_load_tiled_copy, tAcc, tTR_tAcc, tTR_gO, tTR_cO, tTR_rAcc = ( + self._tmem_load_partition( + common_params, common_params.tiled_mma_pv, iter_n + ) + ) + + # load o + cute.copy(tmem_load_tiled_copy, tTR_tAcc, tTR_rAcc) + + # apply output scale and normalize by row_sum + for i in cutlass.range( + cute.size(tTR_rAcc), vectorize=True, unroll_full=True + ): + tTR_rAcc[i] = ( + tTR_rAcc[i] + * epilogue_params.output_scale + * cute.arch.rcp_approx(row_sum) + ) + + # store o to global memory + tR2G_rO_src = None + tR2G_rO_dst = tTR_gO + if cutlass.const_expr(common_params.mAccO is None): + tR2G_rO_src = cute.make_fragment_like(tTR_gO, self.o_dtype) + # using final output dtype for o + tR2G_rO_src.store(tTR_rAcc.load().to(self.o_dtype)) + else: + # using accumulate dtype for o + tR2G_rO_src = tTR_rAcc + + if cute.elem_less(tTR_cO[0][0], common_params.H): + cute.autovec_copy( + tR2G_rO_src, + tR2G_rO_dst, + l1c_evict_priority=cute.nvgpu.CacheEvictionPriority.NO_ALLOCATE, + ) + + # store the lse to global memory + cta_pv_tiler = ( + self.mma_pv_tiler[0] // self.cluster_shape_mnk[0], + self.mma_pv_tiler[1], + self.mma_pv_tiler[2], + ) + gLSE = None + cLSE = None + if cutlass.const_expr(epilogue_params.mAccLSE is None): + gLSE = cute.local_tile( + epilogue_params.mLSE, + (cta_pv_tiler[0], 1, 1), + ( + common_params.blk_coord[0], + common_params.blk_coord[1], + common_params.blk_coord[2], + ), + (1, 1, 1), + ) + cLSE = cute.local_tile( + cute.make_identity_tensor(epilogue_params.mLSE.shape), + (cta_pv_tiler[0], 1, 1), + ( + common_params.blk_coord[0], + common_params.blk_coord[1], + common_params.blk_coord[2], + ), + (1, 1, 1), + ) + + else: + gLSE = cute.local_tile( + epilogue_params.mAccLSE[ + None, common_params.blk_coord[3], None, None + ], + (cta_pv_tiler[0], 1, 1), + ( + common_params.blk_coord[0], + common_params.blk_coord[1], + common_params.blk_coord[2], + ), + (1, 1, 1), + ) + cLSE = cute.local_tile( + cute.make_identity_tensor( + epilogue_params.mAccLSE[ + None, common_params.blk_coord[3], None, None + ].shape + ), + (cta_pv_tiler[0], 1, 1), + ( + common_params.blk_coord[0], + common_params.blk_coord[1], + common_params.blk_coord[2], + ), + (1, 1, 1), + ) + lse = ( + cute.math.log2(row_sum, fastmath=True) + + epilogue_params.softmax_scale_log2 * row_max + ) + if cutlass.const_expr(self.warps_in_n == 2): + if cute.elem_less(cLSE[tidx][0], common_params.H): + gLSE[tidx] = lse + + cute.arch.fence_view_async_tmem_load() + common_params.mma_o_pipeline.consumer_release(mma_o_consumer_state) + mma_o_consumer_state.advance() + + return mma_o_consumer_state + + def make_and_init_load_qkv_pipeline( + self, load_qkv_mbar_ptr, cta_layout_vmnk, load_stages, tx_count + ) -> pipeline.PipelineTmaUmma: + """Create and initialize the tma load qkv pipeline. + + :param load_qkv_mbar_ptr: The load qkv mbar pointer + :type load_qkv_mbar_ptr: cute.Tensor + :param cta_layout_vmnk: The cta layout vmnk + :type cta_layout_vmnk: tuple[int, int, int] + :param load_stages: The load stages + :type load_stages: list[int] + :param tx_count: The tx count + :type tx_count: int + + :return: The tma load qkv pipeline + :rtype: pipeline.PipelineTmaUmma + """ + load_qkv_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, len([self.load_tma_k_warp_id]) + ) + load_qkv_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, len([self.mma_warp_id]) + ) + return pipeline.PipelineTmaUmma.create( + barrier_storage=load_qkv_mbar_ptr, + num_stages=load_stages, + producer_group=load_qkv_producer_group, + consumer_group=load_qkv_consumer_group, + tx_count=tx_count, + cta_layout_vmnk=cta_layout_vmnk, + defer_sync=True, + ) + + def make_and_init_mma_s_pipeline( + self, mma_s_mbar_ptr, cta_layout_vmnk + ) -> pipeline.PipelineUmmaAsync: + """Create and initialize the mma s pipeline. + + :param mma_s_mbar_ptr: The mma s mbar pointer + :type mma_s_mbar_ptr: cute.Tensor + :param cta_layout_vmnk: The cta layout vmnk + :type cta_layout_vmnk: tuple[int, int, int] + + :return: The mma s pipeline + :rtype: pipeline.PipelineUmmaAsync + """ + + mma_s_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, len([self.mma_warp_id]) + ) + consumer_thread_size = ( + self.threads_per_warp + * len(self.compute_warp_ids) + * self.cluster_shape_mnk[0] + ) + mma_s_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + consumer_thread_size, + ) + return pipeline.PipelineUmmaAsync.create( + barrier_storage=mma_s_mbar_ptr, + num_stages=self.mma_s_stage, + producer_group=mma_s_producer_group, + consumer_group=mma_s_consumer_group, + cta_layout_vmnk=cta_layout_vmnk, + defer_sync=True, + ) + + def make_and_init_p_mma_pipeline( + self, p_mma_mbar_ptr, cta_layout_vmnk + ) -> pipeline.PipelineAsyncUmma: + """Create and initialize the p mma pipeline. + + :param p_mma_mbar_ptr: The p mma mbar pointer + :type p_mma_mbar_ptr: cute.Tensor + :param cta_layout_vmnk: The cta layout vmnk + :type cta_layout_vmnk: tuple[int, int, int] + + :return: The p mma pipeline + :rtype: pipeline.PipelineAsyncUmma + """ + + producer_thread_size = ( + self.threads_per_warp + * len(self.compute_warp_ids) + * self.cluster_shape_mnk[0] + ) + p_mma_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + producer_thread_size, + ) + p_mma_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, len([self.mma_warp_id]) + ) + return pipeline.PipelineAsyncUmma.create( + barrier_storage=p_mma_mbar_ptr, + num_stages=self.p_mma_stage, + producer_group=p_mma_producer_group, + consumer_group=p_mma_consumer_group, + cta_layout_vmnk=cta_layout_vmnk, + defer_sync=True, + ) + + def make_and_init_p_cor_pipeline( + self, p_cor_mbar_ptr + ) -> pipeline.PipelineAsyncUmma: + """Create and initialize the p correction pipeline. + + :param p_cor_mbar_ptr: The p correction mbar pointer + :type p_cor_mbar_ptr: cute.Tensor + + :return: The p correction pipeline + :rtype: pipeline.PipelineAsyncUmma + """ + + producer_thread_size = self.threads_per_warp * len(self.compute_warp_ids) + p_cor_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + producer_thread_size, + ) + p_cor_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + producer_thread_size, + ) + return pipeline.PipelineAsync.create( + barrier_storage=p_cor_mbar_ptr, + num_stages=self.p_cor_stage, + producer_group=p_cor_producer_group, + consumer_group=p_cor_consumer_group, + defer_sync=True, + ) + + def make_and_init_mma_o_pipeline( + self, mma_o_mbar_ptr, cta_layout_vmnk + ) -> pipeline.PipelineUmmaAsync: + """Create and initialize the mma o pipeline. + + :param mma_o_mbar_ptr: The mma o mbar pointer + :type mma_o_mbar_ptr: cute.Tensor + :param cta_layout_vmnk: The cta layout vmnk + :type cta_layout_vmnk: tuple[int, int, int] + + :return: The mma o pipeline + :rtype: pipeline.PipelineUmmaAsync + """ + + mma_o_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, len([self.mma_warp_id]) + ) + consumer_thread_size = ( + self.threads_per_warp + * len(self.compute_warp_ids) + * self.cluster_shape_mnk[0] + ) + mma_o_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + consumer_thread_size, + ) + return pipeline.PipelineUmmaAsync.create( + barrier_storage=mma_o_mbar_ptr, + num_stages=self.mma_o_stage, + producer_group=mma_o_producer_group, + consumer_group=mma_o_consumer_group, + cta_layout_vmnk=cta_layout_vmnk, + defer_sync=True, + ) + + @staticmethod + def _compute_grid( + o: cute.Tensor, + split_kv: cutlass.Int32, + cluster_shape_mnk: Tuple[int, int, int], + max_active_clusters: int, + is_persistent: bool, + ) -> Tuple[MLAStaticTileSchedulerParams, Tuple[int, int, int]]: + """Compute grid shape for the output tensor C. + + :param c: The output tensor C + :type c: cute.Tensor + :param cta_tile_shape_mnk: The shape (M, N, K) of the CTA tile. + :type cta_tile_shape_mnk: tuple[int, int, int] + :param cluster_shape_mn: Shape of each cluster in M, N dimensions. + :type cluster_shape_mn: tuple[int, int] + + :return: Tile scheduler parameters and grid shape. + :rtype: tuple[MLAStaticTileSchedulerParams, tuple[int, int, int]] + """ + o_shape = o.shape + tile_sched_params = create_mla_static_tile_scheduler_params( + is_persistent, + cute.size(o_shape[3]), + cute.size(o_shape[2]), + cluster_shape_mnk, + split_kv, + ) + grid = MLAStaticTileScheduler.get_grid_shape( + tile_sched_params, max_active_clusters + ) + + return tile_sched_params, grid + + @staticmethod + def get_workspace_size( + H: int, + S: int, + D: int, + B: int, + split_kv: int, + acc_dtype: Type[cutlass.Numeric], + ) -> int: + """Get the extra workspace(device memory) size for the MLA kernel when split_kv is not 1. + + :param H: The height of the output tensor C + :type H: int + :param S: The sequence length of the output tensor C + :type S: int + :param D: The depth of the output tensor C + :type D: int + :param B: The batch size of the output tensor C + :type B: int + :param split_kv: The split key-value of the output tensor C + :type split_kv: int + :param acc_dtype: The data type of the output tensor C + :type acc_dtype: Type[cutlass.Numeric] + + :return: The workspace size for the MLA kernel + :rtype: int + """ + if split_kv == 1: + return 0 + return B * H * S * split_kv * (D + 1) * acc_dtype.width // 8 + + @cute.jit + def initialize_workspace( + self, + H: cutlass.Int32, + D: cutlass.Int32, + S: cutlass.Int32, + B: cutlass.Int32, + split_kv: cutlass.Int32, + acc_dtype: Type[cutlass.Numeric], + workspace: cute.Tensor, + ) -> tuple[cute.Tensor, cute.Tensor]: + """Initialize the workspace for the MLA kernel. Construct the intermediate tensors + acc_o and acc_lse. + + :param H: The height of the output tensor C + :type H: cutlass.Int32 + :param D: The depth of the output tensor C + :type D: cutlass.Int32 + :param S: The sequence length of the output tensor C + :type S: cutlass.Int32 + :param B: The batch size of the output tensor C + :type B: cutlass.Int32 + :param split_kv: The split key-value of the output tensor C + :type split_kv: cutlass.Int32 + :param acc_dtype: The data type of the output tensor C + :type acc_dtype: Type[cutlass.Numeric] + :param workspace: The workspace tensor + :type workspace: cute.Tensor + + :return: The output tensor C and the workspace tensor + :rtype: tuple[cute.Tensor, cute.Tensor] + """ + acc_o, acc_lse = None, None + if cutlass.const_expr(workspace is not None): + align = 256 // self.q_dtype.width + acc_o_layout = cute.make_layout( + (H, split_kv, D, S, B), + stride=( + cute.assume(split_kv * D, align), + cute.assume(D, align), + 1, + cute.assume(split_kv * H * D, align), + cute.assume(H * split_kv * S * D, align), + ), + ) + acc_o_iter = cute.recast_ptr(workspace.iterator, dtype=acc_dtype) + acc_o = cute.make_tensor(acc_o_iter, acc_o_layout) + acc_lse_layout = cute.make_layout( + (H, split_kv, S, B), + stride=(split_kv, 1, H * split_kv, H * split_kv * S), + ) + acc_lse_iter = cute.recast_ptr( + workspace.iterator + cute.cosize(acc_o_layout) * acc_dtype.width // 8, + dtype=acc_dtype, + ) + acc_lse = cute.make_tensor(acc_lse_iter, acc_lse_layout) + return acc_o, acc_lse + + @staticmethod + def can_implement( + B: int, + S: int, + K: int, + H: int, + L: int, + R: int, + in_dtype: Type[cutlass.Numeric], + out_dtype: Type[cutlass.Numeric], + acc_dtype: Type[cutlass.Numeric], + lse_dtype: Type[cutlass.Numeric], + mma_qk_tiler_mn: Tuple[int, int], + mma_pv_tiler_mn: Tuple[int, int], + split_kv: int, + is_persistent: bool, + is_var_seq: bool, + is_var_split_kv: bool, + page_size: int, + ) -> bool: + """Check if the MLA kernel can be implemented. + + :param B: The batch size of the output tensor C + :type B: int + :param S: The sequence length of the output tensor C + :type S: int + :param K: The width of the output tensor KV + :type K: int + :param H: The number of heads of the output tensor C + :type H: int + :param L: The number of latent dimensions of the tensor KV + :type L: int + :param R: The number of rope dimensions of the tensor C_rope + :type R: int + :param in_dtype: The data type of the input tensor + :type in_dtype: Type[cutlass.Numeric] + :param out_dtype: The data type of the output tensor + :type out_dtype: Type[cutlass.Numeric] + :param acc_dtype: The data type of the accumulator + :type acc_dtype: Type[cutlass.Numeric] + :param lse_dtype: The data type of the log-sum-exp + :type lse_dtype: Type[cutlass.Numeric] + :param mma_qk_tiler_mn: The tile shape of the query-key matrix multiplication + :type mma_qk_tiler_mn: Tuple[int, int] + :param mma_pv_tiler_mn: The tile shape of the probability-value matrix multiplication + :type mma_pv_tiler_mn: Tuple[int, int] + :param split_kv: The split key-value of the output tensor C + :type split_kv: int + :param is_persistent: Whether to use persistent kernel optimization + :type is_persistent: bool + :param is_var_seq: Whether to use variable sequence length + :type is_var_seq: bool + :param is_var_split_kv: Whether to use variable split_kv + :type is_var_split_kv: bool + :param page_size: The page size of the page table + :type page_size: int + + :return: Whether the MLA kernel can be implemented + :rtype: bool + """ + if L != 512 or R != 64: + return False + if in_dtype not in [cutlass.Float8E4M3FN]: + return False + if out_dtype not in [cutlass.Float8E4M3FN]: + return False + if acc_dtype != cutlass.Float32 or lse_dtype != cutlass.Float32: + return False + # page size equals 1 is prohibited by tma specification, not 128B aligned. + if mma_qk_tiler_mn[1] % page_size != 0 or page_size == 1: + return False + if mma_qk_tiler_mn[0] != mma_pv_tiler_mn[0] or mma_qk_tiler_mn[0] != 128: + return False + if is_var_split_kv and not is_var_seq: + return False + if H > 128 or (H < 128 and split_kv != 1): + return False + if S <= 0 or S > 4: + return False + if K <= 0: + return False + return True + + +def run( + batch_size: int, + seq_len_q: int, + seq_len_k: int, + num_heads: int, + latent_dim: int, + rope_dim: int, + in_dtype: Type[cutlass.Numeric], + out_dtype: Type[cutlass.Numeric], + acc_dtype: Type[cutlass.Numeric], + lse_dtype: Type[cutlass.Numeric], + mma_qk_tiler_mn: Tuple[int, int], + mma_pv_tiler_mn: Tuple[int, int], + split_kv: int, + is_persistent: bool, + is_var_seq: bool, + is_var_split_kv: bool, + page_size: int, + softmax_scale: float, + output_scale: float, + skip_correction_threshold: float, + tolerance: float, + warmup_iterations: int, + iterations: int, + skip_ref_check: bool, + use_cold_l2: bool, + **kwargs, +): + """Execute Multi-Head Latent Attention (MLA) on Blackwell architecture and validate results. + + This function creates random input tensors for query latent/rope, compressed latent/rope, and value, + then performs the complete MLA computation pipeline. It supports configurable data types, tiling parameters, + page table, variable sequence length, and variable split_kv. Results can be validated against a PyTorch reference + implementation or run multiple times for performance measurement. + + :param batch_size: Batch size + :type batch_size: int + :param seq_len_q: Sequence length of Q + :type seq_len_q: int + :param seq_len_k: Sequence length of K + :type seq_len_k: int + :param num_heads: Number of heads + :type num_heads: int + :param latent_dim: dimension of query/compressed latent + :type latent_dim: int + :param rope_dim: dimension of query/compressed rope + :type rope_dim: int + :param in_dtype: Input data type for query/compressed latent/rope tensors + :type in_dtype: Type[cutlass.Numeric] + :param out_dtype: Output data type for attention output + :type out_dtype: Type[cutlass.Numeric] + :param acc_dtype: Accumulator data type for query-key matrix multiplication + :type acc_dtype: Type[cutlass.Numeric] + :param lse_dtype: Accumulator data type for log-sum-exp + :type lse_dtype: Type[cutlass.Numeric] + :param mma_qk_tiler_mn: Matrix multiply accumulate tile shape (M, N) for query-key matrix multiplication + :type mma_qk_tiler_mn: Tuple[int, int] + :param mma_pv_tiler_mn: Matrix multiply accumulate tile shape (M, N) for probability-value matrix multiplication + :type mma_pv_tiler_mn: Tuple[int, int] + :param split_kv: Split key-value + :type split_kv: int + :param is_persistent: Whether to use persistent kernel optimization + :type is_persistent: bool + :param is_var_seq: Whether to use variable sequence length + :type is_var_seq: bool + :param is_var_split_kv: Whether to use variable split_kv + :type is_var_split_kv: bool + :param page_size: Page size of the page table + :type page_size: int + :param softmax_scale: Attention score scaling factor + :type softmax_scale: float + :param output_scale: Output scaling factor + :type output_scale: float + :param skip_correction_threshold: Threshold to skip correction + :type skip_correction_threshold: float + :param tolerance: Maximum acceptable error for validation + :type tolerance: float + :param warmup_iterations: Number of warmup iterations + :type warmup_iterations: int + :param iterations: Number of iterations to run for performance testing + :type iterations: int + :param skip_ref_check: Skip validation against reference implementation + :type skip_ref_check: bool + :param use_cold_l2: Whether to use cold L2 cache + :type use_cold_l2: bool + + :raises ValueError: If input shapes are incompatible or head dimension is unsupported + :raises RuntimeError: If GPU is unavailable for computation + """ + + print("Running Blackwell MLA test with:") + print(f" batch_size: {batch_size}") + print(f" seq_len_q: {seq_len_q}") + print(f" seq_len_k: {seq_len_k}") + print(f" num_heads: {num_heads}") + print(f" latent_dim: {latent_dim}") + print(f" rope_dim: {rope_dim}") + print(f" in_dtype: {in_dtype}") + print(f" out_dtype: {out_dtype}") + print(f" acc_dtype: {acc_dtype}") + print(f" mma_qk_tiler_mn: {mma_qk_tiler_mn}") + print(f" mma_pv_tiler_mn: {mma_pv_tiler_mn}") + print(f" split_kv: {split_kv}") + print(f" is_persistent: {is_persistent}") + print(f" is_var_seq: {is_var_seq}") + print(f" is_var_split_kv: {is_var_split_kv}") + print(f" page_size: {page_size}") + print(f" softmax_scale: {softmax_scale}") + print(f" output_scale: {output_scale}") + print(f" skip_correction_threshold: {skip_correction_threshold}") + print(f" tolerance: {tolerance}") + print(f" warmup_iterations: {warmup_iterations}") + print(f" iterations: {iterations}") + print(f" skip_ref_check: {skip_ref_check}") + print(f" use_cold_l2: {use_cold_l2}") + + import torch + import cutlass.torch as cutlass_torch + + # Prepare pytorch tensors: Q, K, V (random from 0 to 2) and O (all zero) + if not torch.cuda.is_available(): + raise RuntimeError("GPU is required to run this example!") + + if not BlackwellMultiHeadLatentAttentionForwardFP8.can_implement( + batch_size, + seq_len_q, + seq_len_k, + num_heads, + latent_dim, + rope_dim, + in_dtype, + out_dtype, + acc_dtype, + lse_dtype, + mma_qk_tiler_mn, + mma_pv_tiler_mn, + split_kv, + is_persistent, + is_var_seq, + is_var_split_kv, + page_size, + ): + raise TypeError( + f"Unsupported testcase {batch_size}, {seq_len_q}, {seq_len_k}, {num_heads}, {latent_dim}, {rope_dim}, {in_dtype}, {out_dtype}, {acc_dtype}, {lse_dtype}, {mma_qk_tiler_mn}, {mma_pv_tiler_mn}, {split_kv}, {is_persistent}, {is_var_seq}, {is_var_split_kv}, {page_size}" + ) + + torch.manual_seed(1111) + + def create_data_tensor( + B, + HK, + D, + dtype, + is_dynamic_layout=True, + page_table=None, + cache_seqs=None, + is_lse=False, + seq_len_q=None, + ): + shape = (B, HK, D) + if page_table is not None: + if cache_seqs is not None: + max_seq_len = torch.max(cache_seqs) + shape = (B * ceil_div(max_seq_len, page_size), page_size, D) + else: + shape = (B * ceil_div(HK, page_size), page_size, D) + + if seq_len_q is not None: + shape = (B, seq_len_q, HK, D) + + permute_order = (1, 2, 0) + stride_order = (2, 0, 1) + leading_dim = 1 + if is_lse: + shape = (B, seq_len_q, HK) + permute_order = (2, 1, 0) + stride_order = (2, 1, 0) + leading_dim = 0 + elif seq_len_q is not None: + permute_order = (2, 3, 1, 0) + stride_order = (3, 2, 0, 1) + leading_dim = 1 + + init_config = cutlass.torch.RandomInitConfig(min_val=-2, max_val=2) + + torch_dtype = ( + cutlass_torch.dtype(dtype) if dtype != cutlass.Float8E4M3FN else torch.int8 + ) + + # Create dtype torch tensor (cpu) + torch_tensor_cpu = cutlass_torch.create_and_permute_torch_tensor( + shape, + torch_dtype, + permute_order=permute_order, + init_type=cutlass.torch.TensorInitType.RANDOM, + init_config=init_config, + ) + + # Create dtype torch tensor (gpu) + torch_tensor_gpu = torch_tensor_cpu.cuda() + + # Create f32 torch tensor (cpu) + f32_torch_tensor = torch_tensor_cpu.to(dtype=torch.float32) + + # Create dtype cute tensor (gpu) + cute_tensor = from_dlpack(torch_tensor_gpu, assumed_align=16) + cute_tensor.element_type = dtype + if is_dynamic_layout: + cute_tensor = cute_tensor.mark_layout_dynamic(leading_dim=leading_dim) + if not is_lse: + cute_tensor = cute_tensor.mark_compact_shape_dynamic( + mode=leading_dim, + stride_order=stride_order, + divisibility=(128 // dtype.width), + ) + + cute_tensor = cutlass_torch.convert_cute_tensor( + f32_torch_tensor, + cute_tensor, + dtype, + is_dynamic_layout=is_dynamic_layout, + ) + + return f32_torch_tensor, cute_tensor, torch_tensor_gpu + + def create_cache_seqs(batch_size, seq_len_k, is_var_seq): + cache_seqs_ref = torch.ones(batch_size, dtype=torch.int32) * seq_len_k + cache_seqs_gpu = cache_seqs_ref.cuda() + cache_seqs = from_dlpack(cache_seqs_gpu, assumed_align=16).mark_layout_dynamic() + if is_var_seq: + max_seq_len = seq_len_k + min_seq_len = int(seq_len_k * 0.8) + cache_seqs_ref = cutlass_torch.create_and_permute_torch_tensor( + (batch_size,), + torch.int32, + init_type=cutlass.torch.TensorInitType.RANDOM, + init_config=cutlass.torch.RandomInitConfig( + min_val=min_seq_len, max_val=max_seq_len + 1 + ), + ) + cache_seqs_gpu = cache_seqs_ref.cuda() + cache_seqs = from_dlpack( + cache_seqs_gpu, + assumed_align=16, + ).mark_layout_dynamic() + return cache_seqs_ref, cache_seqs, cache_seqs_gpu + + def create_page_table(batch_size, seq_len_k, is_var_seq, page_size): + max_seq_len = seq_len_k if not is_var_seq else torch.max(cache_seqs_ref) + page_count = ceil_div(max_seq_len, page_size) + page_table_ref = torch.empty([batch_size, page_count], dtype=torch.int32) + # use transposed index for page table to make sure the value is in bound of `batch_size * seq_len_block`. In practice, the value could be any positive values. This setting is only for testing purpose. + for b in range(batch_size): + for j in range(page_count): + page_table_ref[b, j] = b + j * batch_size + page_table_gpu = page_table_ref.permute(1, 0).cuda() + page_table = from_dlpack(page_table_gpu, assumed_align=16).mark_layout_dynamic( + leading_dim=0 + ) + return page_table_ref, page_table, page_table_gpu + + def create_block_split_kvs( + batch_size, + split_kv, + cache_seqs_ref, + is_var_split_kv, + mma_qk_tiler_mn, + cluster_shape_mnk, + max_active_clusters, + ): + block_split_kvs_ref, block_split_kvs, block_split_kvs_gpu = None, None, None + # check if split_kv is valid otherwise do auto setting of split_kv + if is_var_split_kv: + block_split_kvs_ref = torch.zeros([batch_size], dtype=torch.int32) + for b in range(batch_size): + block_split_kvs_ref[b] = ( + BlackwellMultiHeadLatentAttentionForwardFP8.get_split_kv( + batch_size, + seq_len_q, + cache_seqs_ref[b].item(), + mma_qk_tiler_mn, + max_active_clusters * cluster_shape_mnk[0], + ) + ) + split_kv = torch.max(block_split_kvs_ref).item() + block_split_kvs_gpu = block_split_kvs_ref.cuda() + block_split_kvs = from_dlpack( + block_split_kvs_gpu, assumed_align=16 + ).mark_layout_dynamic() + elif split_kv <= 0: + split_kv = BlackwellMultiHeadLatentAttentionForwardFP8.get_split_kv( + batch_size, + seq_len_q, + cache_seqs_ref[0].item(), + mma_qk_tiler_mn, + max_active_clusters * cluster_shape_mnk[0], + ) + return split_kv, block_split_kvs_ref, block_split_kvs, block_split_kvs_gpu + + def create_workspace( + num_heads, seq_len_q, latent_dim, batch_size, split_kv, acc_dtype + ): + workspace_size = BlackwellMultiHeadLatentAttentionForwardFP8.get_workspace_size( + num_heads, + seq_len_q, + latent_dim, + batch_size, + split_kv, + acc_dtype, + ) + + workspace, workspace_torch = None, None + if workspace_size > 0: + workspace_torch = torch.empty([workspace_size], dtype=torch.int8).cuda() + workspace = from_dlpack(workspace_torch, assumed_align=32) + return workspace, workspace_torch + + cache_seqs_ref, cache_seqs, cache_seqs_torch = create_cache_seqs( + batch_size, seq_len_k, is_var_seq + ) + page_table_ref, page_table, page_table_torch = create_page_table( + batch_size, seq_len_k, is_var_seq, page_size + ) + cluster_shape_mnk = (2, 1, 1) + hardware_info = utils.HardwareInfo() + max_active_clusters = hardware_info.get_max_active_clusters( + cluster_shape_mnk[0] * cluster_shape_mnk[1] + ) + split_kv, block_split_kvs_ref, block_split_kvs, block_split_kvs_torch = ( + create_block_split_kvs( + batch_size, + split_kv, + cache_seqs_ref, + is_var_split_kv, + mma_qk_tiler_mn, + cluster_shape_mnk, + max_active_clusters, + ) + ) + + q_latent_ref, q_latent, q_latent_torch = create_data_tensor( + batch_size, + num_heads, + latent_dim, + in_dtype, + is_dynamic_layout=True, + seq_len_q=seq_len_q, + ) + q_rope_ref, q_rope, q_rope_torch = create_data_tensor( + batch_size, + num_heads, + rope_dim, + in_dtype, + is_dynamic_layout=True, + seq_len_q=seq_len_q, + ) + + c_latent_ref, c_latent, c_latent_torch = create_data_tensor( + batch_size, + seq_len_k, + latent_dim, + in_dtype, + is_dynamic_layout=True, + page_table=page_table, + cache_seqs=cache_seqs_ref, + ) + c_rope_ref, c_rope, c_rope_torch = create_data_tensor( + batch_size, + seq_len_k, + rope_dim, + in_dtype, + is_dynamic_layout=True, + page_table=page_table, + cache_seqs=cache_seqs_ref, + ) + o_ref, o, o_torch = create_data_tensor( + batch_size, + num_heads, + latent_dim, + out_dtype, + is_dynamic_layout=True, + seq_len_q=seq_len_q, + ) + lse_ref, lse, lse_torch = create_data_tensor( + batch_size, + num_heads, + 1, + lse_dtype, + is_dynamic_layout=True, + is_lse=True, + seq_len_q=seq_len_q, + ) + workspace, workspace_torch = create_workspace( + num_heads, seq_len_q, latent_dim, batch_size, split_kv, acc_dtype + ) + + mla = BlackwellMultiHeadLatentAttentionForwardFP8( + acc_dtype, + lse_dtype, + mma_qk_tiler_mn, + mma_pv_tiler_mn, + max_active_clusters, + page_size, + skip_correction_threshold, + is_persistent, + is_var_seq, + is_var_split_kv, + ) + + # Get current CUDA stream from PyTorch + torch_stream = torch.cuda.current_stream() + # Get the raw stream pointer as a CUstream + stream = cuda.CUstream(torch_stream.cuda_stream) + + # compile mla kernel + compiled_mla = cute.compile( + mla, + q_latent, + q_rope, + c_latent, + c_rope, + page_table, + o, + lse, + workspace, + split_kv, + cache_seqs, + block_split_kvs, + softmax_scale, + output_scale, + stream, + options="--opt-level 2", + ) + + def torch_reference_mla( + q_latent, + q_rope, + c_latent, + c_rope, + page_table, + cache_seqs, + softmax_scale=1.0, + output_scale=1.0, + ): + # expand and concat q_latent and q_rope to have the dimension of sequence length for q + q_ref = torch.cat([q_latent, q_rope], dim=1).permute(3, 2, 0, 1) + # expand and concat c_latent and c_rope to have the dimension of num_heads for k and v + page_count = page_table_ref.shape[1] + k_ref_paged = ( + torch.cat([c_latent, c_rope], dim=1) + .permute(2, 0, 1) + .reshape(batch_size * page_count, page_size, latent_dim + rope_dim) + ) + v_ref_paged = c_latent.permute(2, 0, 1).reshape( + batch_size * page_count, page_size, latent_dim + ) + + if is_var_seq: + max_seq_len = torch.max(cache_seqs_ref) + else: + max_seq_len = seq_len_k + + k_ref = torch.zeros([batch_size, 1, max_seq_len, latent_dim + rope_dim]) + v_ref = torch.zeros([batch_size, 1, max_seq_len, latent_dim]) + k_ref = torch.index_select( + k_ref_paged, 0, torch.flatten(page_table_ref) + ).reshape(batch_size, 1, -1, latent_dim + rope_dim)[:, :, :max_seq_len, :] + v_ref = torch.index_select( + v_ref_paged, 0, torch.flatten(page_table_ref) + ).reshape(batch_size, 1, -1, latent_dim)[:, :, :max_seq_len, :] + for b in range(batch_size): + k_ref[b, :, cache_seqs_ref[b] :, :] = 0 + v_ref[b, :, cache_seqs_ref[b] :, :] = 0 + import torch.nn.functional as F + + o_ref = F.scaled_dot_product_attention( + q_ref, + k_ref, + v_ref, + attn_mask=None, + dropout_p=0.0, + scale=softmax_scale, + is_causal=False, + ) + s_ref = torch.einsum("bhld,bhsd->bhls", q_ref, k_ref) + s_ref_max, s_ref_max_pos = torch.max(s_ref, dim=-1, keepdim=True) + softmax_scale_log2 = LOG2_E * softmax_scale + s_ref_sum = torch.sum( + torch.exp2((s_ref - s_ref_max) * softmax_scale_log2), dim=-1, keepdim=True + ) + + lse_ref = s_ref_max * softmax_scale_log2 + torch.log2(s_ref_sum) + lse_ref = lse_ref.squeeze(3).permute(2, 1, 0) + o_ref = o_ref * output_scale + o_ref = o_ref.permute(2, 3, 1, 0) + + return o_ref, lse_ref + + if skip_correction_threshold > 0.0: + print( + "Skipping correction verification since skip_correction_threshold is greater than 0.0..." + ) + skip_ref_check = True + if not skip_ref_check: + # Execute kernel once for reference checking + compiled_mla( + q_latent, + q_rope, + c_latent, + c_rope, + page_table, + o, + lse, + workspace, + split_kv, + cache_seqs, + block_split_kvs, + softmax_scale, + output_scale, + stream, + ) + torch.cuda.synchronize() + + print("Verifying results...") + if in_dtype == cutlass.Float8E4M3FN: + tolerance = 0.13 + o_ref, lse_ref = torch_reference_mla( + q_latent_ref, + q_rope_ref, + c_latent_ref, + c_rope_ref, + page_table, + cache_seqs, + softmax_scale, + output_scale, + ) + + if out_dtype in [cutlass.Float8E5M2, cutlass.Float8E4M3FN]: + # convert o back to f32 for comparison + o_fp32, o_fp32_torch = cutlass_torch.cute_tensor_like( + torch.empty(*o_torch.shape, dtype=torch.float32), + cutlass.Float32, + is_dynamic_layout=True, + assumed_align=16, + ) + cute.testing.convert(o, o_fp32) + o = o_fp32_torch.cpu() + ref_fp8, _ = cutlass_torch.cute_tensor_like( + torch.empty( + *o_ref.permute(3, 2, 0, 1).shape, dtype=torch.uint8 + ).permute(2, 3, 1, 0), + out_dtype, + is_dynamic_layout=True, + assumed_align=16, + ) + o_ref_gpu = o_ref.cuda() + o_ref_f32 = from_dlpack(o_ref_gpu).mark_layout_dynamic(leading_dim=1) + + # convert ref : f32 -> fp8 -> f32 + cute.testing.convert(o_ref_f32, ref_fp8) + cute.testing.convert(ref_fp8, o_ref_f32) + + o_ref = o_ref_gpu.cpu() + else: + o = o_torch.cpu().to(torch.float32) + lse = lse_torch.cpu() + lse_ref = lse_ref.to(cutlass.torch.dtype(lse_dtype)) + # Assert close results + torch.testing.assert_close(o, o_ref, atol=tolerance, rtol=1e-05) + torch.testing.assert_close(lse, lse_ref, atol=tolerance, rtol=1e-05) + print("Results verified successfully!") + + def generate_tensors(): + _, cache_seqs, _ = create_cache_seqs(batch_size, seq_len_k, is_var_seq) + _, page_table, _ = create_page_table( + batch_size, seq_len_k, is_var_seq, page_size + ) + _split_kv, _, block_split_kvs, _ = create_block_split_kvs( + batch_size, + split_kv, + cache_seqs_ref, + is_var_split_kv, + mma_qk_tiler_mn, + cluster_shape_mnk, + max_active_clusters, + ) + + _, q_latent, _ = create_data_tensor( + batch_size, + num_heads, + latent_dim, + in_dtype, + is_dynamic_layout=True, + seq_len_q=seq_len_q, + ) + _, q_rope, _ = create_data_tensor( + batch_size, + num_heads, + rope_dim, + in_dtype, + is_dynamic_layout=True, + seq_len_q=seq_len_q, + ) + + _, c_latent, _ = create_data_tensor( + batch_size, + seq_len_k, + latent_dim, + in_dtype, + is_dynamic_layout=True, + page_table=page_table, + cache_seqs=cache_seqs_ref, + ) + _, c_rope, _ = create_data_tensor( + batch_size, + seq_len_k, + rope_dim, + in_dtype, + is_dynamic_layout=True, + page_table=page_table, + cache_seqs=cache_seqs_ref, + ) + _, o, _ = create_data_tensor( + batch_size, + num_heads, + latent_dim, + out_dtype, + is_dynamic_layout=True, + seq_len_q=seq_len_q, + ) + _, lse, _ = create_data_tensor( + batch_size, + num_heads, + 1, + lse_dtype, + is_dynamic_layout=True, + is_lse=True, + seq_len_q=seq_len_q, + ) + workspace, workspace_torch = create_workspace( + num_heads, seq_len_q, latent_dim, batch_size, _split_kv, acc_dtype + ) + return testing.JitArguments( + q_latent, + q_rope, + c_latent, + c_rope, + page_table, + o, + lse, + workspace, + _split_kv, + cache_seqs, + block_split_kvs, + softmax_scale, + output_scale, + stream, + ) + + workspace_count = 1 + if use_cold_l2: + one_workspace_bytes = ( + q_latent_torch.numel() * q_latent_torch.element_size() + + q_rope_torch.numel() * q_rope_torch.element_size() + + c_latent_torch.numel() * c_latent_torch.element_size() + + c_rope_torch.numel() * c_rope_torch.element_size() + + o_torch.numel() * o_torch.element_size() + + lse_torch.numel() * lse_torch.element_size() + + cache_seqs_torch.numel() * cache_seqs_torch.element_size() + ) + one_workspace_bytes += ( + page_table_torch.numel() * page_table_torch.element_size() + ) + if is_var_split_kv: + one_workspace_bytes += ( + block_split_kvs_torch.numel() * block_split_kvs_torch.element_size() + ) + if workspace_torch is not None: + one_workspace_bytes += ( + workspace_torch.numel() * workspace_torch.element_size() + ) + workspace_count = testing.get_workspace_count( + one_workspace_bytes, warmup_iterations, iterations + ) + + avg_time_us = testing.benchmark( + compiled_mla, + workspace_generator=generate_tensors, + workspace_count=workspace_count, + stream=stream, + warmup_iterations=warmup_iterations, + iterations=iterations, + ) + + return avg_time_us # Return execution time in microseconds + + +if __name__ == "__main__": + + def parse_comma_separated_ints(s: str) -> Tuple[int, ...]: + try: + return tuple(int(x.strip()) for x in s.split(",")) + except ValueError: + raise argparse.ArgumentTypeError( + "Invalid format. Expected comma-separated integers." + ) + + def parse_mma_tiler(s: str) -> Tuple[int, int, Tuple[int, int]]: + ret = parse_comma_separated_ints(s) + if len(ret) != 2: + raise argparse.ArgumentTypeError( + "Invalid format. Expected 2 comma-separated integers." + ) + return (ret[0], ret[1]) + + parser = argparse.ArgumentParser(description="Example of MLA on Blackwell.") + + parser.add_argument( + "--in_dtype", + type=cutlass.dtype, + default=cutlass.Float8E4M3FN, + help="Input data type", + ) + + parser.add_argument( + "--out_dtype", + type=cutlass.dtype, + default=cutlass.Float8E4M3FN, + help="Output data type", + ) + + parser.add_argument( + "--acc_dtype", + type=cutlass.dtype, + default=cutlass.Float32, + help="Accumulator data type", + ) + + parser.add_argument( + "--lse_dtype", + type=cutlass.dtype, + default=cutlass.Float32, + help="LSE data type", + ) + parser.add_argument( + "--mma_qk_tiler_mn", + type=parse_mma_tiler, + default=(128, 128), + help="MMA tile shape (H, K)", + ) + parser.add_argument( + "--mma_pv_tiler_mn", + type=parse_mma_tiler, + default=(128, 256), + help="MMA tile shape (H, D)", + ) + + parser.add_argument( + "--is_persistent", + action="store_true", + help="Is persistent", + ) + + parser.add_argument( + "--batch_size", + type=int, + default=1, + help="Batch size", + ) + + parser.add_argument( + "--seq_len_q", + type=int, + default=1, + help="Sequence length of Q", + ) + + parser.add_argument( + "--seq_len_k", + type=int, + default=128, + help="Sequence length of K/V", + ) + + parser.add_argument( + "--num_heads", + type=int, + default=128, + help="Number of heads of Q", + ) + + parser.add_argument( + "--latent_dim", + type=int, + default=512, + help="Latent dimension of Q/C", + ) + + parser.add_argument( + "--rope_dim", + type=int, + default=64, + help="Rope dimension of Q/C", + ) + + parser.add_argument( + "--is_var_seq", + action="store_true", + help="Use variable length of sequence length or not", + ) + + parser.add_argument( + "--is_var_split_kv", + action="store_true", + help="Use variable length of split kv or not", + ) + + parser.add_argument( + "--page_size", + type=int, + default=128, + help="Page size of page table", + ) + + parser.add_argument( + "--split_kv", + type=int, + default=-1, + help="Split KV setting", + ) + + parser.add_argument( + "--softmax_scale", + type=float, + default=0.0416, + help="Scaling factor to scale softmax", + ) + + parser.add_argument( + "--output_scale", + type=float, + default=1.0, + help="Scaling factor to scale output", + ) + parser.add_argument( + "--skip_correction_threshold", + type=float, + default=0.0, + help="Threshold to skip correction", + ) + + parser.add_argument( + "--tolerance", type=float, default=1e-02, help="Tolerance for validation" + ) + + parser.add_argument( + "--warmup_iterations", + type=int, + default=0, + help="Number of iterations for warmup", + ) + + parser.add_argument( + "--iterations", + type=int, + default=1, + help="Number of iterations after warmup", + ) + + parser.add_argument( + "--skip_ref_check", + action="store_true", + help="Skip reference check", + ) + + parser.add_argument( + "--use_cold_l2", + action="store_true", + help="Use cold L2 cache", + ) + + args = parser.parse_args() + + run( + args.batch_size, + args.seq_len_q, + args.seq_len_k, + args.num_heads, + args.latent_dim, + args.rope_dim, + args.in_dtype, + args.out_dtype, + args.acc_dtype, + args.lse_dtype, + args.mma_qk_tiler_mn, + args.mma_pv_tiler_mn, + args.split_kv, + args.is_persistent, + args.is_var_seq, + args.is_var_split_kv, + args.page_size, + args.softmax_scale, + args.output_scale, + args.skip_correction_threshold, + args.tolerance, + args.warmup_iterations, + args.iterations, + args.skip_ref_check, + args.use_cold_l2, + ) + + print("PASS") diff --git a/examples/python/CuTeDSL/blackwell/mla/mla_helpers.py b/examples/python/CuTeDSL/blackwell/mla/mla_helpers.py new file mode 100644 index 00000000..1790b3c8 --- /dev/null +++ b/examples/python/CuTeDSL/blackwell/mla/mla_helpers.py @@ -0,0 +1,304 @@ +# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + + +import cutlass +import cutlass.cute as cute + + +class MLAStaticTileSchedulerParams: + def __init__( + self, + is_persistent: bool, + problem_shape_b: cute.Int32, + problem_shape_s: cute.Int32, + cluster_shape_mnk: cute.Shape, + split_kv: cutlass.Int32, + *, + problem_shape_b_fdd: cute.FastDivmodDivisor = None, + problem_shape_s_fdd: cute.FastDivmodDivisor = None, + split_kv_fdd: cute.FastDivmodDivisor = None, + loc=None, + ip=None, + ): + """The static tile scheduler parameters prepared for MLA static tile scheduler. + + :param is_persistent: Whether to use persistent kernel mode + :type is_persistent: bool + :param problem_shape_b: The shape of the problem + :type problem_shape_b: cute.Int32 + :param problem_shape_s: The shape of the problem in sequence length Q dimension + :type problem_shape_s: cute.Int32 + :param cluster_shape_mnk: The shape of the cluster + :type cluster_shape_mnk: cute.Shape + :param split_kv: The scalar factor for split KV + """ + self.is_persistent = is_persistent + self.problem_shape_b = problem_shape_b + self.problem_shape_s = problem_shape_s + self.problem_shape_b_fdd = problem_shape_b_fdd + self.problem_shape_s_fdd = problem_shape_s_fdd + self.cluster_shape_mnk = cluster_shape_mnk + self.split_kv = split_kv + self.split_kv_fdd = split_kv_fdd + if cutlass.const_expr(problem_shape_b_fdd is None): + self.problem_shape_b_fdd = cute.fast_divmod_create_divisor( + problem_shape_b, loc=loc, ip=ip + ) + if cutlass.const_expr(problem_shape_s_fdd is None): + self.problem_shape_s_fdd = cute.fast_divmod_create_divisor( + problem_shape_s, loc=loc, ip=ip + ) + if cutlass.const_expr(split_kv_fdd is None): + self.split_kv_fdd = cute.fast_divmod_create_divisor( + split_kv, loc=loc, ip=ip + ) + self.loc = loc + self.ip = ip + + def __extract_mlir_values__(self): + values = cutlass.extract_mlir_values(self.problem_shape_b) + values += cutlass.extract_mlir_values(self.problem_shape_s) + values += cutlass.extract_mlir_values(self.split_kv) + values += cutlass.extract_mlir_values(self.problem_shape_b_fdd) + values += cutlass.extract_mlir_values(self.problem_shape_s_fdd) + values += cutlass.extract_mlir_values(self.split_kv_fdd) + return values + + def __new_from_mlir_values__(self, values): + problem_shape_b = cutlass.new_from_mlir_values( + self.problem_shape_b, (values[0],) + ) + problem_shape_s = cutlass.new_from_mlir_values( + self.problem_shape_s, (values[1],) + ) + split_kv = cutlass.new_from_mlir_values(self.split_kv, (values[2],)) + problem_shape_b_fdd = cutlass.new_from_mlir_values( + self.problem_shape_b_fdd, (values[3],) + ) + problem_shape_s_fdd = cutlass.new_from_mlir_values( + self.problem_shape_s_fdd, (values[4],) + ) + split_kv_fdd = cutlass.new_from_mlir_values(self.split_kv_fdd, (values[5],)) + return MLAStaticTileSchedulerParams( + self.is_persistent, + problem_shape_b, + problem_shape_s, + self.cluster_shape_mnk, + split_kv, + problem_shape_b_fdd=problem_shape_b_fdd, + problem_shape_s_fdd=problem_shape_s_fdd, + split_kv_fdd=split_kv_fdd, + loc=self.loc, + ) + + +def create_mla_static_tile_scheduler_params( + is_persistent: bool, + problem_shape_b: cute.Int32, + problem_shape_s: cute.Int32, + cluster_shape_mnk: cute.Shape, + split_kv: cutlass.Int32, +) -> MLAStaticTileSchedulerParams: + return MLAStaticTileSchedulerParams( + is_persistent, problem_shape_b, problem_shape_s, cluster_shape_mnk, split_kv + ) + + +class WorkTileInfo: + def __init__(self, blk_coord: cute.Coord, is_valid: bool): + self.blk_coord = blk_coord + self.is_valid = cutlass.Boolean(is_valid) + + def __extract_mlir_values__(self): + values = cutlass.extract_mlir_values(self.blk_coord) + values += cutlass.extract_mlir_values(self.is_valid) + return values + + def __new_from_mlir_values__(self, values): + new_tile_idx = cutlass.new_from_mlir_values(self.blk_coord, values[:-1]) + new_is_valid_tile = cutlass.new_from_mlir_values(self.is_valid, [values[-1]]) + return WorkTileInfo(new_tile_idx, new_is_valid_tile) + + @property + def is_valid_tile(self) -> cutlass.Boolean: + return self.is_valid + + @property + def tile_idx(self) -> cute.Coord: + return self.blk_coord + + +class MLAStaticTileScheduler: + def __init__( + self, + params: MLAStaticTileSchedulerParams, + current_work_linear_idx: cutlass.Int32, + blk_coord: cute.Coord, + grid_shape: cute.Shape, + *, + is_valid: bool = True, + loc=None, + ip=None, + ): + """The static tile scheduler for MLA split kv kernel. + Based on `is_persistent`, it provides 2 modes for use: + - Persistent mode: Launch fixed blocks and reschedule the data blocks. + - Non-persistent mode: Launch dynamic blocks and exit when the current work is done. + + :param params: The static tile scheduler parameters + :type params: MLAStaticTileSchedulerParams + :param current_work_linear_idx: The linear index of the current work + :type current_work_linear_idx: cutlass.Int32 + :param blk_coord: The coordinate of the current work + :type blk_coord: cute.Coord + :param grid_shape: The shape of the grid + :type grid_shape: cute.Shape + :param is_valid: Whether the current work is valid + :type is_valid: bool + """ + self.params = params + self.blk_coord = blk_coord + self.grid_shape = grid_shape + self.current_work_linear_idx = current_work_linear_idx + if params.is_persistent: + self.persistent_blk_layout = cute.make_layout( + ( + params.cluster_shape_mnk[0], + params.problem_shape_s, + params.problem_shape_b, + params.split_kv, + ), + loc=loc, + ip=ip, + ) + self.num_blocks = cute.size(self.persistent_blk_layout, loc=loc, ip=ip) + # Used for persistent scheduling + self.num_persistent_sm = cute.size(grid_shape, loc=loc, ip=ip) + else: + self.is_valid = is_valid + self.loc = loc + self.ip = ip + + @staticmethod + def get_grid_shape( + params: MLAStaticTileSchedulerParams, + max_active_clusters: int, + *, + loc=None, + ip=None, + ) -> cute.Shape: + # called by host + grid_shape = ( + params.cluster_shape_mnk[0], + params.problem_shape_b * params.problem_shape_s, + params.split_kv, + ) + if params.is_persistent: + return ( + cutlass.min( + max_active_clusters * cute.size(params.cluster_shape_mnk), + cute.size(grid_shape, loc=loc, ip=ip), + ), + 1, + 1, + ) + else: + return grid_shape + + def get_current_work(self, *, loc=None, ip=None) -> WorkTileInfo: + is_valid = ( + self.current_work_linear_idx < self.num_blocks + if self.params.is_persistent + else self.is_valid + ) + + if self.params.is_persistent: + current_work_cluster_batch, cluster_idx = ( + self.current_work_linear_idx // self.params.cluster_shape_mnk[0], + self.current_work_linear_idx % self.params.cluster_shape_mnk[0], + ) + current_work_s_batch, s_idx = divmod( + current_work_cluster_batch, self.params.problem_shape_s_fdd + ) + current_work_b_batch, b_idx = divmod( + current_work_s_batch, self.params.problem_shape_b_fdd + ) + _, split_kv_idx = divmod(current_work_b_batch, self.params.split_kv_fdd) + + blk_coord = (cluster_idx, s_idx, b_idx, split_kv_idx) + else: + s_idx, b_idx = divmod(self.blk_coord[1], self.params.problem_shape_b_fdd) + blk_coord = (self.blk_coord[0], s_idx, b_idx, self.blk_coord[2]) + + return WorkTileInfo(blk_coord, is_valid) + + def initial_work_tile_info(self, *, loc=None, ip=None): + return self.get_current_work(loc=loc, ip=ip) + + def advance_to_next_work(self, *, advance_count=1, loc=None, ip=None): + if self.params.is_persistent: + self.current_work_linear_idx += advance_count * self.num_persistent_sm + else: + self.is_valid = False + + def __extract_mlir_values__(self): + values = cutlass.extract_mlir_values(self.params) + values.extend(cutlass.extract_mlir_values(self.current_work_linear_idx)) + values.extend(cutlass.extract_mlir_values(self.blk_coord)) + values.extend(cutlass.extract_mlir_values(self.grid_shape)) + return values + + def __new_from_mlir_values__(self, values): + assert len(values) == 13 + new_params = cutlass.new_from_mlir_values(self.params, values[0:6]) + new_current_work_linear_idx = cutlass.new_from_mlir_values( + self.current_work_linear_idx, [values[6]] + ) + new_blk_coord = cutlass.new_from_mlir_values(self.blk_coord, values[7:10]) + new_grid_shape = cutlass.new_from_mlir_values(self.grid_shape, values[10:]) + return MLAStaticTileScheduler( + new_params, new_current_work_linear_idx, new_blk_coord, new_grid_shape + ) + + +def create_mla_static_tile_scheduler( + params: MLAStaticTileSchedulerParams, + blk_coord: cute.Coord, + grid_shape: cute.Shape, +) -> MLAStaticTileScheduler: + return MLAStaticTileScheduler(params, blk_coord[0], blk_coord, grid_shape) + + +LOG2_E = 1.4426950408889634074 +# avoid register indexing on array. +MAX_SPLITS = 256 + + +def ceil_div(a: int, b: int) -> int: + return (a + b - 1) // b diff --git a/examples/python/CuTeDSL/blackwell/sm103_dense_blockscaled_gemm_persistent.py b/examples/python/CuTeDSL/blackwell/sm103_dense_blockscaled_gemm_persistent.py index 607a4e4f..39f14848 100644 --- a/examples/python/CuTeDSL/blackwell/sm103_dense_blockscaled_gemm_persistent.py +++ b/examples/python/CuTeDSL/blackwell/sm103_dense_blockscaled_gemm_persistent.py @@ -304,7 +304,12 @@ class Sm103BlockScaledPersistentDenseGemmKernel: self.is_sfb_mcast = self.num_mcast_ctas_sfb > 1 # Compute epilogue subtile - self.epi_tile = (self.cta_tile_shape_mnk[0], 64) + self.epi_tile = sm103_utils.compute_epilogue_tile_shape( + self.cta_tile_shape_mnk, + self.use_2cta_instrs, + self.c_layout, + self.c_dtype, + ) self.num_acc_stage, self.num_ab_stage, self.num_sf_stage, self.num_c_stage = ( self._compute_stages( @@ -371,6 +376,38 @@ class Sm103BlockScaledPersistentDenseGemmKernel: self.c_dtype, self.c_layout, self.epi_tile, self.num_c_stage ) + # Overlap and double buffer accumulator when num_acc_stage == 1 for cta_tile_n = 256 case + self.overlapping_accum = self.num_acc_stage == 1 and not self.use_tma_store + self.epi_tile_n = cute.size(self.epi_tile[1]) + + if self.overlapping_accum: + # Compute SF TMEM column count from a scale factor layout. + # Column count = cosize of Int32-recast layout & 0xFFFF, + # mirroring the computation in find_tmem_tensor_col_offset. + def _sf_tmem_cols(make_tmem_layout_fn, smem_layout_staged): + layout = make_tmem_layout_fn( + tiled_mma, + self.mma_tiler, + self.sf_vec_size, + cute.slice_(smem_layout_staged, (None, None, None, 0)), + ) + return ( + cute.cosize(cute.recast_layout(32, self.sf_dtype.width, layout)) + & 0xFFFF + ) + + self.num_sfa_tmem_cols = _sf_tmem_cols( + blockscaled_utils.make_tmem_layout_sfa, self.sfa_smem_layout_staged + ) + self.num_sfb_tmem_cols = _sf_tmem_cols( + blockscaled_utils.make_tmem_layout_sfb, self.sfb_smem_layout_staged + ) + self.num_sf_tmem_cols = self.num_sfa_tmem_cols + self.num_sfb_tmem_cols + # Release accumulator buffer early in epilogue when overlapping + self.iter_acc_early_release_in_epilogue = ( + self.num_sf_tmem_cols // self.epi_tile_n + ) + @cute.jit def __call__( self, @@ -944,10 +981,30 @@ class Sm103BlockScaledPersistentDenseGemmKernel: # # (MMA, MMA_M, MMA_N) acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2]) - # (MMA, MMA_M, MMA_N, STAGE) - tCtAcc_fake = tiled_mma.make_fragment_C( - cute.append(acc_shape, self.num_acc_stage) - ) + if cutlass.const_expr(self.overlapping_accum): + num_acc_stage_overlapped = 2 + tCtAcc_fake = tiled_mma.make_fragment_C( + cute.append(acc_shape, num_acc_stage_overlapped) + ) + # (MMA, MMA_M, MMA_N, STAGE) + tCtAcc_fake = cute.make_tensor( + tCtAcc_fake.iterator, + cute.make_layout( + tCtAcc_fake.shape, + stride=( + tCtAcc_fake.stride[0], + tCtAcc_fake.stride[1], + tCtAcc_fake.stride[2], + (self.cta_tile_shape_mnk[1] - self.num_sf_tmem_cols) + * tCtAcc_fake.stride[0][1], + ), + ), + ) + else: + # (MMA, MMA_M, MMA_N, STAGE) + tCtAcc_fake = tiled_mma.make_fragment_C( + cute.append(acc_shape, self.num_acc_stage) + ) # # Cluster wait before tensor memory alloc @@ -1244,8 +1301,14 @@ class Sm103BlockScaledPersistentDenseGemmKernel: cur_tile_coord[2], ) + # Get accumulator stage index + if cutlass.const_expr(self.overlapping_accum): + acc_stage_index = acc_producer_state.phase ^ 1 + else: + acc_stage_index = acc_producer_state.index + # Set tensor memory buffer for current tile - tCtAcc = tCtAcc_base[(None, 0, 0, acc_producer_state.index)] + tCtAcc = tCtAcc_base[(None, 0, 0, acc_stage_index)] # Peek (try_wait) AB buffer full for k_tile = 0 ab_consumer.reset() @@ -1665,6 +1728,7 @@ class Sm103BlockScaledPersistentDenseGemmKernel: acc_pipeline, tCcC_base=tCcC, mC_mnl=mC_mnl, + overlapping_accum=self.overlapping_accum, ) if cutlass.const_expr(self.use_tma_store): @@ -1680,8 +1744,6 @@ class Sm103BlockScaledPersistentDenseGemmKernel: tmem.relinquish_alloc_permit() tmem.free(acc_tmem_ptr) - cute.arch.mbarrier_init_fence() - @staticmethod def make_desc_and_call_mma( tiled_mma: cute.TiledMma, diff --git a/examples/python/CuTeDSL/blackwell/tutorial_gemm/fp16_gemm_0.py b/examples/python/CuTeDSL/blackwell/tutorial_gemm/fp16_gemm_0.py index 341b47d2..6f05b56e 100644 --- a/examples/python/CuTeDSL/blackwell/tutorial_gemm/fp16_gemm_0.py +++ b/examples/python/CuTeDSL/blackwell/tutorial_gemm/fp16_gemm_0.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2024 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: LicenseRef-NvidiaProprietary # # NVIDIA CORPORATION, its affiliates and licensors retain all intellectual @@ -9,13 +9,11 @@ # its affiliates is strictly prohibited. import argparse -import torch from typing import Tuple import cutlass import cutlass.cute as cute import cutlass.utils as utils -import cutlass.torch as cutlass_torch import cutlass.pipeline as pipeline from cutlass.cute.nvgpu import cpasync, tcgen05 import cutlass.utils.blackwell_helpers as sm100_utils @@ -32,9 +30,8 @@ with optimizations for challenges that may arise with other problem sizes. To run this example: .. code-block:: bash - python examples/blackwell/tutorial_fp16_gemm_0.py \ - --mnk 8192,8192,8192 \ - --tolerance 1e-01 + python examples/blackwell/tutorial_gemm/fp16_gemm_0.py \ + --mnk 8192,8192,8192 Constraints for this example: * The problem size of m and n must be divisible by the tile size m & n (128, 256) @@ -128,7 +125,8 @@ def kernel( num_stages=acc_stage, producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), consumer_group=pipeline.CooperativeGroup( - pipeline.Agent.Thread, threads_per_cta + pipeline.Agent.Thread, + threads_per_cta, ), barrier_storage=storage.acc_mbar_ptr.data_ptr(), ).make_participants() @@ -141,15 +139,15 @@ def kernel( # (bM, bN) gC = cute.local_tile(mC_mnl, mma_tiler_mnk, mma_coord_mnk, proj=(1, 1, None)) thr_mma = tiled_mma.get_slice(0) - # (MMA, MMA_M, MMA_K, RestK) + # (MMA, MMA_M, MMA_K) tCgA = thr_mma.partition_A(gA) - # (MMA, MMA_N, MMA_K, RestK) + # (MMA, MMA_N, MMA_K) tCgB = thr_mma.partition_B(gB) # (MMA, MMA_M, MMA_N) tCgC = thr_mma.partition_C(gC) - # (MMA, MMA_M, MMA_K, STAGE) + # (MMA, MMA_M, MMA_K) tCrA = tiled_mma.make_fragment_A(sA) - # (MMA, MMA_N, MMA_K, STAGE) + # (MMA, MMA_N, MMA_K) tCrB = tiled_mma.make_fragment_B(sB) # (MMA, MMA_M, MMA_N) acc_shape = tiled_mma.partition_shape_C(mma_tiler_mnk[:2]) @@ -188,7 +186,7 @@ def kernel( # (EpiTile, NumTiles) gC_epi = cute.zipped_divide(tCgC, epi_tiler) - # Every thread loads 32x128 bits + # Every thread loads 64 x fp32 tmem_atom = cute.make_copy_atom( tcgen05.Ld32x32bOp(tcgen05.Repetition.x64), cutlass.Float32, @@ -273,11 +271,7 @@ def kernel( @cute.jit -def host_function( - a: cute.Tensor, - b: cute.Tensor, - c: cute.Tensor, -): +def host_function(a: cute.Tensor, b: cute.Tensor, c: cute.Tensor): # Construct tiled MMA op = tcgen05.MmaF16BF16Op( io_dtype, @@ -354,6 +348,10 @@ def run_dense_gemm( mnk: Tuple[int, int, int], tolerance: float, ): + global torch, cutlass_torch + import torch + import cutlass.torch as cutlass_torch + print("===================================================================") print("Running Blackwell fp16 GEMM example 0 with:") print(f" mnk: {mnk}") @@ -393,12 +391,7 @@ def run_dense_gemm( ) # Entry point to the host JIT function - host_function( - a_tensor, - b_tensor, - c_tensor, - no_cache=True, - ) + host_function(a_tensor, b_tensor, c_tensor, no_cache=True) # Compute reference result and verify ref = (torch.einsum("mk,nk->mn", a.to(torch.float32), b.to(torch.float32))).cpu() @@ -418,7 +411,11 @@ if __name__ == "__main__": "Invalid format. Expected comma-separated integers." ) - if not torch.cuda.is_available(): + from cuda.bindings import driver as cu_driver + + cu_driver.cuInit(0) + err, device_count = cu_driver.cuDeviceGetCount() + if err != cu_driver.CUresult.CUDA_SUCCESS or device_count < 1: raise RuntimeError("A GPU is required to run this example") parser = argparse.ArgumentParser(description="Blackwell fp16 GEMM example 0") diff --git a/examples/python/CuTeDSL/blackwell/tutorial_gemm/fp16_gemm_1.py b/examples/python/CuTeDSL/blackwell/tutorial_gemm/fp16_gemm_1.py index 1d1d98ae..12e74fa0 100644 --- a/examples/python/CuTeDSL/blackwell/tutorial_gemm/fp16_gemm_1.py +++ b/examples/python/CuTeDSL/blackwell/tutorial_gemm/fp16_gemm_1.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2024 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: LicenseRef-NvidiaProprietary # # NVIDIA CORPORATION, its affiliates and licensors retain all intellectual @@ -13,15 +13,14 @@ import argparse -import torch from typing import Tuple import cutlass import cutlass.cute as cute import cutlass.utils as utils -import cutlass.torch as cutlass_torch import cutlass.pipeline as pipeline from cutlass.cute.nvgpu import cpasync, tcgen05 +import cutlass.utils.blackwell_helpers as sm100_utils from cutlass.cute.runtime import from_dlpack """ @@ -78,7 +77,6 @@ acc_stage = 1 @cute.struct class SharedStorage: - # each stage has 2 kinds of barrier, i.e. empty & full ab_mbar_ptr: cute.struct.MemRange[cutlass.Int64, ab_stages * 2] acc_mbar_ptr: cute.struct.MemRange[cutlass.Int64, acc_stage * 2] tmem_dealloc_mbar_ptr: cutlass.Int64 @@ -174,15 +172,15 @@ def kernel( # (bM, bN) gC = cute.local_tile(mC_mnl, mma_tiler_mnk, mma_coord_mnk, proj=(1, 1, None)) thr_mma = tiled_mma.get_slice(mma_coord_vmnk[0]) - # (MMA, MMA_M, MMA_K, RestK) + # (MMA, MMA_M, MMA_K) tCgA = thr_mma.partition_A(gA) - # (MMA, MMA_N, MMA_K, RestK) + # (MMA, MMA_N, MMA_K) tCgB = thr_mma.partition_B(gB) # (MMA, MMA_M, MMA_N) tCgC = thr_mma.partition_C(gC) - # (MMA, MMA_M, MMA_K, STAGE) + # (MMA, MMA_M, MMA_K) tCrA = tiled_mma.make_fragment_A(sA) - # (MMA, MMA_N, MMA_K, STAGE) + # (MMA, MMA_N, MMA_K) tCrB = tiled_mma.make_fragment_B(sB) # (MMA, MMA_M, MMA_N) acc_shape = tiled_mma.partition_shape_C(mma_tiler_mnk[:2]) @@ -256,9 +254,9 @@ def kernel( tDgC = tmem_thr_copy.partition_D(gC_epi) # (TmemCpy,NumTmemCpy) - tCrAcc = cute.make_rmem_tensor_like(tDgC[None, None, 0], acc_dtype) + tCrAcc = cute.make_rmem_tensor(tDgC[None, None, 0].shape, acc_dtype) # (TmemCpy,NumTmemCpy) - tCrC = cute.make_rmem_tensor_like(tDgC[None, None, 0], io_dtype) + tCrC = cute.make_rmem_tensor(tDgC[None, None, 0].shape, io_dtype) # # 2. Main loop @@ -356,13 +354,13 @@ def host_function( tiled_mma = cute.make_tiled_mma(op) # Construct SMEM layouts for A and B - a_smem_layout = utils.sm100.make_smem_layout_a( + a_smem_layout = sm100_utils.make_smem_layout_a( tiled_mma, mma_tiler_mnk, a.element_type, ab_stages, ) - b_smem_layout = utils.sm100.make_smem_layout_b( + b_smem_layout = sm100_utils.make_smem_layout_b( tiled_mma, mma_tiler_mnk, b.element_type, @@ -383,7 +381,7 @@ def host_function( a_smem_layout_one_stage, mma_tiler_mnk, tiled_mma, - cta_layout_vmnk.shape, + cta_layout_vmnk.shape, # take the layout and extract the shape internally ) b_tma_atom, b_tma_tensor = cute.nvgpu.make_tiled_tma_atom_B( op, @@ -438,6 +436,10 @@ def run_dense_gemm( mnk: Tuple[int, int, int], tolerance: float, ): + global torch, cutlass_torch + import torch + import cutlass.torch as cutlass_torch + print("===================================================================") print("Running Blackwell fp16 GEMM example 1 with:") print(f" mnk: {mnk}") @@ -501,7 +503,11 @@ if __name__ == "__main__": "Invalid format. Expected comma-separated integers." ) - if not torch.cuda.is_available(): + from cuda.bindings import driver as cu_driver + + cu_driver.cuInit(0) + err, device_count = cu_driver.cuDeviceGetCount() + if err != cu_driver.CUresult.CUDA_SUCCESS or device_count < 1: raise RuntimeError("A GPU is required to run this example") parser = argparse.ArgumentParser(description="Blackwell fp16 GEMM example 1") diff --git a/examples/python/CuTeDSL/blackwell/tutorial_gemm/nvfp4_gemm_0.py b/examples/python/CuTeDSL/blackwell/tutorial_gemm/nvfp4_gemm_0.py new file mode 100644 index 00000000..f0dfeea1 --- /dev/null +++ b/examples/python/CuTeDSL/blackwell/tutorial_gemm/nvfp4_gemm_0.py @@ -0,0 +1,778 @@ +# 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. + + +import argparse +import os +import sys +from typing import Type, Tuple +import cuda.bindings.driver as cuda + +import torch + +import cutlass +import cutlass.cute as cute +import cutlass.utils as utils +import cutlass.pipeline as pipeline +from cutlass.cute.nvgpu import cpasync, tcgen05 +import cutlass.torch as cutlass_torch +import cutlass.utils.blackwell_helpers as sm100_utils +import cutlass.utils.blockscaled_layout as blockscaled_utils +from cutlass.cute.runtime import make_ptr + +if __name__ == "__main__": + current_dir = os.path.dirname(os.path.abspath(__file__)) + examples_dir = os.path.join(current_dir, "..", "..") + if examples_dir not in sys.path: + sys.path.insert(0, examples_dir) + +from blackwell.tutorial_gemm.utils import create_parser, run + +mma_tiler_mn = (128, 256) +mma_inst_shape_k = 64 +ab_dtype = cutlass.Float4E2M1FN +sf_dtype = cutlass.Float8E4M3FN +c_dtype = cutlass.Float16 +sf_vec_size = 16 + +""" +The first tutorial NVFP4 block-scaled batched GEMM demonstrating a simple kernel implementation in CuTeDSL + +This example demonstrates the kernel implementation of block-scaled batched GEMM with NVFP4 data type. +With large tile sizes (128x256x256), it can achieve very high performance on 8k×8k×8k problem sizes. +It can serve as a starting point to help users quickly experiment with optimizations for +challenges that may arise with other problem sizes. + +To run this example: +.. code-block:: bash + + python examples/blackwell/tutorial_gemm/nvfp4_gemm_0.py \ + --mnkl 8192,8192,8192,1 --do_benchmark + +Constraints for this example: +* The problem size of m, n and k must be divisible by the tile size m&n&k (128,256,256) +* The scaling factor vector size is 16. +* The A/B matrices have data contiguous on the k dimension. +* The C matrix has data contiguous on the n dimension. +* The A/B matrix data type is Float4E2M1FN. +* The SFA/SFB matrix data type is Float8E4M3FN. +""" + + +class Sm100BlockScaledDenseGemmKernel: + def __init__(self): + self.threads_per_cta = 128 + self.smem_capacity = utils.get_smem_capacity_in_bytes("sm_100") + self.num_tmem_alloc_cols = 512 + + # set stages for ab_pipeline and acc_pipeline + self.num_acc_stage = 1 + self.num_ab_stage = 4 + + @cute.jit + def __call__( + self, + a_ptr: cute.Pointer, + b_ptr: cute.Pointer, + sfa_ptr: cute.Pointer, + sfb_ptr: cute.Pointer, + c_ptr: cute.Pointer, + problem_size: tuple, + stream: cuda.CUstream, + epilogue_op: cutlass.Constexpr = lambda x: x, + ): + # setup static attributes before smem/grid/tma computation + self.c_layout = utils.LayoutEnum.ROW_MAJOR + m, n, k, l = problem_size + + # Setup attributes that depend on gemm inputs + mma_inst_tile_k = 4 + self.mma_tiler = ( + mma_tiler_mn[0], + mma_tiler_mn[1], + mma_inst_shape_k * mma_inst_tile_k, + ) + self.cta_tile_shape_mnk = ( + self.mma_tiler[0], + self.mma_tiler[1], + self.mma_tiler[2], + ) + + a_tensor = cute.make_tensor( + a_ptr, + cute.make_layout( + (m, cute.assume(k, 32), l), + stride=(cute.assume(k, 32), 1, cute.assume(m * k, 32)), + ), + ) + b_tensor = cute.make_tensor( + b_ptr, + cute.make_layout( + (n, cute.assume(k, 32), l), + stride=(cute.assume(k, 32), 1, cute.assume(n * k, 32)), + ), + ) + # make address offset of c_tensor 256bit aligned, + # so that epilogue could use vectorized store with larger vector size. + c_tensor = cute.make_tensor( + c_ptr, + cute.make_layout( + (cute.assume(m, 32), cute.assume(n, 16), l), + stride=(cute.assume(n, 16), 1, cute.assume(m * n, 512)), + ), + ) + # Setup sfa/sfb tensor by filling A/B tensor to scale factor atom layout + # ((Atom_M, Rest_M),(Atom_K, Rest_K),RestL) + sfa_layout = blockscaled_utils.tile_atom_to_shape_SF( + a_tensor.shape, sf_vec_size + ) + sfa_tensor = cute.make_tensor(sfa_ptr, sfa_layout) + + # ((Atom_N, Rest_N),(Atom_K, Rest_K),RestL) + sfb_layout = blockscaled_utils.tile_atom_to_shape_SF( + b_tensor.shape, sf_vec_size + ) + sfb_tensor = cute.make_tensor(sfb_ptr, sfb_layout) + + mma_op = tcgen05.MmaMXF4NVF4Op( + sf_dtype, + (*mma_tiler_mn, mma_inst_shape_k), + tcgen05.CtaGroup.ONE, + tcgen05.OperandSource.SMEM, + ) + tiled_mma = cute.make_tiled_mma(mma_op) + + self.cluster_layout_vmnk = cute.tiled_divide( + cute.make_layout((1, 1, 1)), + (tiled_mma.thr_id.shape,), + ) + + # Compute A/B/SFA/SFB/C shared memory layout + self.a_smem_layout_staged = sm100_utils.make_smem_layout_a( + tiled_mma, + self.mma_tiler, + ab_dtype, + self.num_ab_stage, + ) + self.b_smem_layout_staged = sm100_utils.make_smem_layout_b( + tiled_mma, + self.mma_tiler, + ab_dtype, + self.num_ab_stage, + ) + self.sfa_smem_layout_staged = blockscaled_utils.make_smem_layout_sfa( + tiled_mma, + self.mma_tiler, + sf_vec_size, + self.num_ab_stage, + ) + self.sfb_smem_layout_staged = blockscaled_utils.make_smem_layout_sfb( + tiled_mma, + self.mma_tiler, + sf_vec_size, + self.num_ab_stage, + ) + + atom_thr_size = cute.size(tiled_mma.thr_id.shape) + + # TMA load for A + a_smem_layout = cute.slice_(self.a_smem_layout_staged, (None, None, None, 0)) + tma_atom_a, tma_tensor_a = cute.nvgpu.make_tiled_tma_atom_A( + cpasync.CopyBulkTensorTileG2SOp(tcgen05.CtaGroup.ONE), + a_tensor, + a_smem_layout, + self.mma_tiler, + tiled_mma, + self.cluster_layout_vmnk.shape, + ) + # TMA load for B + b_smem_layout = cute.slice_(self.b_smem_layout_staged, (None, None, None, 0)) + tma_atom_b, tma_tensor_b = cute.nvgpu.make_tiled_tma_atom_B( + cpasync.CopyBulkTensorTileG2SOp(tcgen05.CtaGroup.ONE), + b_tensor, + b_smem_layout, + self.mma_tiler, + tiled_mma, + self.cluster_layout_vmnk.shape, + ) + + # TMA load for SFA + sfa_smem_layout = cute.slice_( + self.sfa_smem_layout_staged, (None, None, None, 0) + ) + tma_atom_sfa, tma_tensor_sfa = cute.nvgpu.make_tiled_tma_atom_A( + cpasync.CopyBulkTensorTileG2SOp(tcgen05.CtaGroup.ONE), + sfa_tensor, + sfa_smem_layout, + self.mma_tiler, + tiled_mma, + self.cluster_layout_vmnk.shape, + internal_type=cutlass.Int16, + ) + + # TMA load for SFB + sfb_smem_layout = cute.slice_( + self.sfb_smem_layout_staged, (None, None, None, 0) + ) + tma_atom_sfb, tma_tensor_sfb = cute.nvgpu.make_tiled_tma_atom_B( + cpasync.CopyBulkTensorTileG2SOp(tcgen05.CtaGroup.ONE), + sfb_tensor, + sfb_smem_layout, + self.mma_tiler, + tiled_mma, + self.cluster_layout_vmnk.shape, + internal_type=cutlass.Int16, + ) + + # Compute TMA load bytes + a_copy_size = cute.size_in_bytes(ab_dtype, a_smem_layout) + b_copy_size = cute.size_in_bytes(ab_dtype, b_smem_layout) + sfa_copy_size = cute.size_in_bytes(sf_dtype, sfa_smem_layout) + sfb_copy_size = cute.size_in_bytes(sf_dtype, sfb_smem_layout) + self.num_tma_load_bytes = ( + a_copy_size + b_copy_size + sfa_copy_size + sfb_copy_size + ) * atom_thr_size + + # Compute grid size + grid = ( + cute.ceil_div(c_tensor.shape[0], self.cta_tile_shape_mnk[0]), + cute.ceil_div(c_tensor.shape[1], self.cta_tile_shape_mnk[1]), + c_tensor.shape[2], + ) + + # Launch the kernel synchronously + self.kernel( + tiled_mma, + tma_atom_a, + tma_tensor_a, + tma_atom_b, + tma_tensor_b, + tma_atom_sfa, + tma_tensor_sfa, + tma_atom_sfb, + tma_tensor_sfb, + c_tensor, + self.a_smem_layout_staged, + self.b_smem_layout_staged, + self.sfa_smem_layout_staged, + self.sfb_smem_layout_staged, + epilogue_op, + ).launch( + grid=grid, + block=[self.threads_per_cta, 1, 1], + cluster=(1, 1, 1), + stream=stream, + ) + return + + # GPU device kernel + @cute.kernel + def kernel( + self, + tiled_mma: cute.TiledMma, + tma_atom_a: cute.CopyAtom, + mA_mkl: cute.Tensor, + tma_atom_b: cute.CopyAtom, + mB_nkl: cute.Tensor, + tma_atom_sfa: cute.CopyAtom, + mSFA_mkl: cute.Tensor, + tma_atom_sfb: cute.CopyAtom, + mSFB_nkl: cute.Tensor, + mC_mnl: cute.Tensor, + a_smem_layout_staged: cute.ComposedLayout, + b_smem_layout_staged: cute.ComposedLayout, + sfa_smem_layout_staged: cute.Layout, + sfb_smem_layout_staged: cute.Layout, + epilogue_op: cutlass.Constexpr, + ): + """ + GPU device kernel performing the batched GEMM computation. + """ + warp_idx = cute.arch.warp_idx() + warp_idx = cute.arch.make_warp_uniform(warp_idx) + tidx, _, _ = cute.arch.thread_idx() + + # + # Setup cta/thread coordinates + # + # Coords inside cluster + bidx, bidy, bidz = cute.arch.block_idx() + + # Coords outside cluster + cta_coord = (bidx, bidy, bidz) + mma_tile_coord_mnl = ( + cta_coord[0] // cute.size(tiled_mma.thr_id.shape), + cta_coord[1], + cta_coord[2], + ) + + # + # Define shared storage for kernel + # + @cute.struct + class SharedStorage: + ab_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] + acc_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] + tmem_holding_buf: cutlass.Int32 + + smem = utils.SmemAllocator() + storage = smem.allocate(SharedStorage) + # (MMA, MMA_M, MMA_K, STAGE) + sA = smem.allocate_tensor( + element_type=ab_dtype, + layout=a_smem_layout_staged.outer, + byte_alignment=128, + swizzle=a_smem_layout_staged.inner, + ) + # (MMA, MMA_N, MMA_K, STAGE) + sB = smem.allocate_tensor( + element_type=ab_dtype, + layout=b_smem_layout_staged.outer, + byte_alignment=128, + swizzle=b_smem_layout_staged.inner, + ) + # (MMA, MMA_M, MMA_K, STAGE) + sSFA = smem.allocate_tensor( + element_type=sf_dtype, + layout=sfa_smem_layout_staged, + byte_alignment=128, + ) + # (MMA, MMA_N, MMA_K, STAGE) + sSFB = smem.allocate_tensor( + element_type=sf_dtype, + layout=sfb_smem_layout_staged, + byte_alignment=128, + ) + + # + # Initialize mainloop ab_pipeline, acc_pipeline and their states + # + ab_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) + ab_pipeline_consumer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, 1) + ab_producer, ab_consumer = pipeline.PipelineTmaUmma.create( + barrier_storage=storage.ab_mbar_ptr.data_ptr(), + num_stages=self.num_ab_stage, + producer_group=ab_pipeline_producer_group, + consumer_group=ab_pipeline_consumer_group, + tx_count=self.num_tma_load_bytes, + ).make_participants() + acc_producer, acc_consumer = pipeline.PipelineUmmaAsync.create( + barrier_storage=storage.acc_mbar_ptr.data_ptr(), + num_stages=self.num_acc_stage, + producer_group=ab_pipeline_producer_group, + consumer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, + self.threads_per_cta, + ), + ).make_participants() + + # + # Local_tile partition global tensors + # + # (bM, bK, RestM, RestK, RestL) + gA_mkl = cute.local_tile( + mA_mkl, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None) + ) + # (bN, bK, RestN, RestK, RestL) + gB_nkl = cute.local_tile( + mB_nkl, cute.slice_(self.mma_tiler, (0, None, None)), (None, None, None) + ) + gSFA_mkl = cute.local_tile( + mSFA_mkl, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None) + ) + gSFB_nkl = cute.local_tile( + mSFB_nkl, cute.slice_(self.mma_tiler, (0, None, None)), (None, None, None) + ) + # (bM, bN, RestM, RestN, RestL) + gC_mnl = cute.local_tile( + mC_mnl, cute.slice_(self.mma_tiler, (None, None, 0)), (None, None, None) + ) + k_tile_cnt = cute.size(gA_mkl, mode=[3]) + + # + # Partition global tensor for TiledMMA_A/B/SFA/SFB/C + # + # (MMA, MMA_M, MMA_K, RestK) + thr_mma = tiled_mma.get_slice(0) + # (MMA, MMA_M, MMA_K, RestM, RestK, RestL) + tCgA = thr_mma.partition_A(gA_mkl) + # (MMA, MMA_N, MMA_K, RestN, RestK, RestL) + tCgB = thr_mma.partition_B(gB_nkl) + # (MMA, MMA_M, MMA_K, RestM, RestK, RestL) + tCgSFA = thr_mma.partition_A(gSFA_mkl) + # (MMA, MMA_N, MMA_K, RestN, RestK, RestL) + tCgSFB = thr_mma.partition_B(gSFB_nkl) + # (MMA, MMA_M, MMA_N, RestM, RestN, RestL) + tCgC = thr_mma.partition_C(gC_mnl) + + # + # Partition global/shared tensor for TMA load A/B/SFA/SFB + # + # TMA load A partition_S/D + # ((atom_v, rest_v), STAGE) + # ((atom_v, rest_v), RestM, RestK, RestL) + tAsA, tAgA = cpasync.tma_partition( + tma_atom_a, + 0, + cute.make_layout(1), + cute.group_modes(sA, 0, 3), + cute.group_modes(tCgA, 0, 3), + ) + # TMA load B partition_S/D + # ((atom_v, rest_v), STAGE) + # ((atom_v, rest_v), RestN, RestK, RestL) + tBsB, tBgB = cpasync.tma_partition( + tma_atom_b, + 0, + cute.make_layout(1), + cute.group_modes(sB, 0, 3), + cute.group_modes(tCgB, 0, 3), + ) + + # TMA load partition for SFA tensor + # ((atom_v, rest_v), STAGE) + # ((atom_v, rest_v), RestM, RestK, RestL) + tAsSFA, tAgSFA = cpasync.tma_partition( + tma_atom_sfa, + 0, + cute.make_layout(1), + cute.group_modes(sSFA, 0, 3), + cute.group_modes(tCgSFA, 0, 3), + ) + tAsSFA = cute.filter_zeros(tAsSFA) + tAgSFA = cute.filter_zeros(tAgSFA) + + # TMA load partition for SFB tensor + # ((atom_v, rest_v), STAGE) + # ((atom_v, rest_v), RestN, RestK, RestL) + tBsSFB, tBgSFB = cpasync.tma_partition( + tma_atom_sfb, + 0, + cute.make_layout(1), + cute.group_modes(sSFB, 0, 3), + cute.group_modes(tCgSFB, 0, 3), + ) + tBsSFB = cute.filter_zeros(tBsSFB) + tBgSFB = cute.filter_zeros(tBgSFB) + + # + # Partition shared/tensor memory tensor for TiledMMA_A/B/C + # + # (MMA, MMA_M, MMA_K, STAGE) + tCrA = tiled_mma.make_fragment_A(sA) + # (MMA, MMA_N, MMA_K, STAGE) + tCrB = tiled_mma.make_fragment_B(sB) + # (MMA, MMA_M, MMA_N) + acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2]) + # (MMA, MMA_M, MMA_N) + tCtAcc_fake = tiled_mma.make_fragment_C(acc_shape) + + # + # Alloc tensor memory buffer + # + tmem_alloc_barrier = pipeline.NamedBarrier( + barrier_id=1, + num_threads=self.threads_per_cta, + ) + tmem = utils.TmemAllocator( + storage.tmem_holding_buf, + barrier_for_retrieve=tmem_alloc_barrier, + ) + tmem.allocate(self.num_tmem_alloc_cols) + tmem.wait_for_alloc() + acc_tmem_ptr = tmem.retrieve_ptr(cutlass.Float32) + tCtAcc = cute.make_tensor(acc_tmem_ptr, tCtAcc_fake.layout) + + # + # Make SFA/SFB tmem tensor + # + # Get SFA tmem ptr + sfa_tmem_ptr = cute.recast_ptr( + acc_tmem_ptr + tcgen05.find_tmem_tensor_col_offset(tCtAcc), + dtype=sf_dtype, + ) + # (MMA, MMA_M, MMA_K) + tCtSFA_layout = blockscaled_utils.make_tmem_layout_sfa( + tiled_mma, + self.mma_tiler, + sf_vec_size, + cute.slice_(sfa_smem_layout_staged, (None, None, None, 0)), + ) + tCtSFA = cute.make_tensor(sfa_tmem_ptr, tCtSFA_layout) + # Get SFB tmem ptr + sfb_tmem_ptr = cute.recast_ptr( + acc_tmem_ptr + + tcgen05.find_tmem_tensor_col_offset(tCtAcc) + + tcgen05.find_tmem_tensor_col_offset(tCtSFA), + dtype=sf_dtype, + ) + # (MMA, MMA_N, MMA_K) + tCtSFB_layout = blockscaled_utils.make_tmem_layout_sfb( + tiled_mma, + self.mma_tiler, + sf_vec_size, + cute.slice_(sfb_smem_layout_staged, (None, None, None, 0)), + ) + tCtSFB = cute.make_tensor(sfb_tmem_ptr, tCtSFB_layout) + + # + # Partition for S2T copy of SFA/SFB + # + # Make S2T CopyAtom + copy_atom_s2t = cute.make_copy_atom( + tcgen05.Cp4x32x128bOp(tcgen05.CtaGroup.ONE), + sf_dtype, + ) + # (MMA, MMA_MN, MMA_K, STAGE) + tCsSFA_compact = cute.filter_zeros(sSFA) + # (MMA, MMA_MN, MMA_K) + tCtSFA_compact = cute.filter_zeros(tCtSFA) + tiled_copy_s2t_sfa = tcgen05.make_s2t_copy(copy_atom_s2t, tCtSFA_compact) + thr_copy_s2t_sfa = tiled_copy_s2t_sfa.get_slice(0) + # ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K, STAGE) + tCsSFA_compact_s2t_ = thr_copy_s2t_sfa.partition_S(tCsSFA_compact) + # ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K, STAGE) + tCsSFA_compact_s2t = tcgen05.get_s2t_smem_desc_tensor( + tiled_copy_s2t_sfa, tCsSFA_compact_s2t_ + ) + # ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K) + tCtSFA_compact_s2t = thr_copy_s2t_sfa.partition_D(tCtSFA_compact) + + # (MMA, MMA_MN, MMA_K, STAGE) + tCsSFB_compact = cute.filter_zeros(sSFB) + # (MMA, MMA_MN, MMA_K) + tCtSFB_compact = cute.filter_zeros(tCtSFB) + tiled_copy_s2t_sfb = tcgen05.make_s2t_copy(copy_atom_s2t, tCtSFB_compact) + thr_copy_s2t_sfb = tiled_copy_s2t_sfb.get_slice(0) + # ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K, STAGE) + tCsSFB_compact_s2t_ = thr_copy_s2t_sfb.partition_S(tCsSFB_compact) + # ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K, STAGE) + tCsSFB_compact_s2t = tcgen05.get_s2t_smem_desc_tensor( + tiled_copy_s2t_sfb, tCsSFB_compact_s2t_ + ) + # ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K) + tCtSFB_compact_s2t = thr_copy_s2t_sfb.partition_D(tCtSFB_compact) + + # + # Slice to per mma tile index + # + # ((atom_v, rest_v), RestK) + tAgA = tAgA[(None, mma_tile_coord_mnl[0], None, mma_tile_coord_mnl[2])] + # ((atom_v, rest_v), RestK) + tBgB = tBgB[(None, mma_tile_coord_mnl[1], None, mma_tile_coord_mnl[2])] + # ((atom_v, rest_v), RestK) + tAgSFA = tAgSFA[(None, mma_tile_coord_mnl[0], None, mma_tile_coord_mnl[2])] + # ((atom_v, rest_v), RestK) + tBgSFB = tBgSFB[(None, mma_tile_coord_mnl[1], None, mma_tile_coord_mnl[2])] + + # + # Execute Data copy and Math computation in the k_tile loop + # + if warp_idx == 0: + # Wait for accumulator buffer empty + acc_empty = acc_producer.acquire_and_advance() + # Set ACCUMULATE field to False for the first k_tile iteration + tiled_mma.set(tcgen05.Field.ACCUMULATE, False) + # Execute k_tile loop + for k_tile in cutlass.range( + k_tile_cnt, prefetch_stages=self.num_ab_stage - 2 + ): + # Wait for AB buffer empty + ab_empty = ab_producer.acquire_and_advance() + + # TMA load for A/B/SFA/SFB + cute.copy( + tma_atom_a, + tAgA[(None, ab_empty.count)], + tAsA[(None, ab_empty.index)], + tma_bar_ptr=ab_empty.barrier, + ) + cute.copy( + tma_atom_b, + tBgB[(None, ab_empty.count)], + tBsB[(None, ab_empty.index)], + tma_bar_ptr=ab_empty.barrier, + ) + cute.copy( + tma_atom_sfa, + tAgSFA[(None, ab_empty.count)], + tAsSFA[(None, ab_empty.index)], + tma_bar_ptr=ab_empty.barrier, + ) + cute.copy( + tma_atom_sfb, + tBgSFB[(None, ab_empty.count)], + tBsSFB[(None, ab_empty.index)], + tma_bar_ptr=ab_empty.barrier, + ) + + # Wait for AB buffer full + ab_full = ab_consumer.wait_and_advance() + + # Copy SFA/SFB to tmem + s2t_stage_coord = (None, None, None, None, ab_full.index) + tCsSFA_compact_s2t_staged = tCsSFA_compact_s2t[s2t_stage_coord] + tCsSFB_compact_s2t_staged = tCsSFB_compact_s2t[s2t_stage_coord] + cute.copy( + tiled_copy_s2t_sfa, + tCsSFA_compact_s2t_staged, + tCtSFA_compact_s2t, + ) + cute.copy( + tiled_copy_s2t_sfb, + tCsSFB_compact_s2t_staged, + tCtSFB_compact_s2t, + ) + + # tCtAcc += tCrA * tCrSFA * tCrB * tCrSFB + num_kblocks = cute.size(tCrA, mode=[2]) + for kblock_idx in cutlass.range(num_kblocks, unroll_full=True): + kblock_coord = ( + None, + None, + kblock_idx, + ab_full.index, + ) + + # Set SFA/SFB tensor to tiled_mma + sf_kblock_coord = (None, None, kblock_idx) + tiled_mma.set( + tcgen05.Field.SFA, + tCtSFA[sf_kblock_coord].iterator, + ) + tiled_mma.set( + tcgen05.Field.SFB, + tCtSFB[sf_kblock_coord].iterator, + ) + + cute.gemm( + tiled_mma, + tCtAcc, + tCrA[kblock_coord], + tCrB[kblock_coord], + tCtAcc, + ) + # Enable accumulate on tCtAcc after first kblock + tiled_mma.set(tcgen05.Field.ACCUMULATE, True) + + # Async arrive AB buffer empty + ab_full.release() + acc_empty.commit() + + # + # Epilogue + # Partition for epilogue + # + op = tcgen05.Ld32x32bOp(tcgen05.Repetition.x128, tcgen05.Pack.NONE) + copy_atom_t2r = cute.make_copy_atom(op, cutlass.Float32) + tiled_copy_t2r = tcgen05.make_tmem_copy(copy_atom_t2r, tCtAcc) + thr_copy_t2r = tiled_copy_t2r.get_slice(tidx) + # (T2R_M, T2R_N, EPI_M, EPI_M) + tTR_tAcc = thr_copy_t2r.partition_S(tCtAcc) + # (T2R_M, T2R_N, EPI_M, EPI_N, RestM, RestN, RestL) + tTR_gC = thr_copy_t2r.partition_D(tCgC) + # (T2R_M, T2R_N, EPI_M, EPI_N） + tTR_rAcc = cute.make_rmem_tensor( + tTR_gC[None, None, None, None, 0, 0, 0].shape, cutlass.Float32 + ) + # (T2R_M, T2R_N, EPI_M, EPI_N） + tTR_rC = cute.make_rmem_tensor( + tTR_gC[None, None, None, None, 0, 0, 0].shape, c_dtype + ) + # STG Atom + simt_atom = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), c_dtype) + tTR_gC = tTR_gC[(None, None, None, None, *mma_tile_coord_mnl)] + + # Release TMEM allocation lock + tmem.relinquish_alloc_permit() + + # Wait for accumulator buffer full + acc_full = acc_consumer.wait_and_advance() + + # Copy accumulator to register + cute.copy(tiled_copy_t2r, tTR_tAcc, tTR_rAcc) + acc_vec = epilogue_op(tTR_rAcc.load().to(c_dtype)) + tTR_rC.store(acc_vec) + # Store C to global memory + cute.copy(simt_atom, tTR_rC, tTR_gC) + + acc_full.release() + + # Deallocate TMEM + cute.arch.barrier() + tmem.free(acc_tmem_ptr) + + return + + +def run_nvfp4_gemm( + mnkl: Tuple[int, int, int, int], + tolerance: float, + do_benchmark: bool = False, + warmup_iterations: int = 10, + iterations: int = 100, + use_cold_l2: bool = True, +): + run( + gemm_class=Sm100BlockScaledDenseGemmKernel, + ab_dtype=ab_dtype, + sf_dtype=sf_dtype, + c_dtype=c_dtype, + sf_vec_size=sf_vec_size, + mma_tiler_mn=mma_tiler_mn, + cluster_shape_mnk=(1, 1, 1), + mnkl=mnkl, + tolerance=tolerance, + do_benchmark=do_benchmark, + warmup_iterations=warmup_iterations, + iterations=iterations, + use_cold_l2=use_cold_l2, + ) + + +if __name__ == "__main__": + parser = create_parser() + args = parser.parse_args() + + if len(args.mnkl) != 4: + parser.error("--mnkl must contain exactly 4 values") + + m, n, k, _ = args.mnkl + if m % mma_tiler_mn[0] != 0: + parser.error("m must be multiples of mma_tiler_mn[0] (got m={})".format(m)) + if n % mma_tiler_mn[1] != 0: + parser.error("n must be multiples of mma_tiler_mn[1] (got n={})".format(n)) + if k % 256 != 0: + parser.error("k must be a multiple of 256 (got k={})".format(k)) + + run_nvfp4_gemm( + args.mnkl, + args.tolerance, + args.do_benchmark, + ) + print("PASS") diff --git a/examples/python/CuTeDSL/blackwell/tutorial_gemm/nvfp4_gemm_1.py b/examples/python/CuTeDSL/blackwell/tutorial_gemm/nvfp4_gemm_1.py new file mode 100644 index 00000000..715de1b7 --- /dev/null +++ b/examples/python/CuTeDSL/blackwell/tutorial_gemm/nvfp4_gemm_1.py @@ -0,0 +1,934 @@ +# 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. + +# This is the second tutorial nvfp4 GEMM. It builds on the first tutorial by adding 2CTA MMA +# instructions with a 2x1 cluster. + +import argparse +import os +import sys +from typing import Type, Tuple +import cuda.bindings.driver as cuda + +import torch + +import cutlass +import cutlass.cute as cute +import cutlass.utils as utils +import cutlass.pipeline as pipeline +from cutlass.cute.nvgpu import cpasync, tcgen05 +import cutlass.torch as cutlass_torch +import cutlass.utils.blackwell_helpers as sm100_utils +import cutlass.utils.blockscaled_layout as blockscaled_utils +from cutlass.cute.runtime import from_dlpack, make_ptr + +if __name__ == "__main__": + current_dir = os.path.dirname(os.path.abspath(__file__)) + examples_dir = os.path.join(current_dir, "..", "..") + if examples_dir not in sys.path: + sys.path.insert(0, examples_dir) + +from blackwell.tutorial_gemm.utils import create_parser, run + +mma_tiler_mn = (256, 256) +mma_inst_shape_k = 64 +ab_dtype = cutlass.Float4E2M1FN +sf_dtype = cutlass.Float8E4M3FN +c_dtype = cutlass.Float16 +sf_vec_size = 16 +cluster_shape_mnk = (2, 1, 1) + +""" +The second tutorial further improves the performance of NVFP4 block-scaled batched GEMM +by adding 2CTA instructions and TMA multicast optimizations. + +(1) The 2 CTA instructions could reduce the smem size requirement for B tensor, +increased num_ab_stage and improves the latency hiding capability. + +For both 1CTA and 2CTA, the shared memory (smem) size per stage for the A, sfA, and sfB tensors is the same: +- For the A tensor, each stage requires 128 x 256 x sizeof(float4) = 16KB. +- For the sfA tensor, each stage requires 128 x (256 / 16) x sizeof(float8) = 2KB. +- For the sfB tensor, each stage requires 256 x (256 / 16) x sizeof(float8) = 4KB. + +The situation is different for the B tensor: +- In the 1CTA case, each stage for the B tensor requires 256 x 256 x sizeof(float4) = 32KB. +- In the 2CTA case, only half this size is needed, i.e., 128 x 256 x sizeof(float4) = 16KB. + +Therefore, the maximum number of AB stages is: +- For 1CTA: 227 // (16 + 32 + 2 + 4) = 4 +- For 2CTA: 227 // (16 + 16 + 2 + 4) = 5 + +The latency hiding capability is: +- 1CTA: 512 * (4 - 1) = 1.5K cycles +- 2CTA: 512 * (5 - 1) = 2K cycles + +(2) TMA multicast can help reduce L2 cache traffic. + +Without TMA multicast, the L2 traffic per tile is typically 16KB + 32KB = 48KB (possibly less in practice, depending on hardware optimizations). +With TMA multicast in a cluster of shape (m, n), the L2 traffic per tile is reduced to 16KB / n + 32KB / m. +For example: +- In a 2x1 cluster: 16KB / 1 + 32KB / 2 = 24KB per tile +- In a 4x4 cluster: 16KB / 4 + 32KB / 4 = 12KB per tile + +The first approach offers substantial capacity for hiding latency, whereas the second reduces the time required for data to become ready. +Both could be tried when the workload is latency-bound or limited by memory throughput. + +To run this example: +.. code-block:: bash + + python examples/blackwell/tutorial_gemm/nvfp4_gemm_1.py \ + --mnkl 8192,8192,8192,1 --do_benchmark + +Constraints for this example: +* The problem size of m, n and k must be divisible by the tile size m&n&k (256, 256, 256) +* The scaling factor vector size is 16. +* The A/B matrices have data contiguous on the k dimension. +* The C matrix has data contiguous on the n dimension. +* The A/B matrix data type is Float4E2M1FN. +* The SFA/SFB matrix data type is Float8E4M3FN. +""" + + +class Sm100BlockScaledDenseGemmKernel: + def __init__(self): + self.threads_per_cta = 128 + self.smem_capacity = utils.get_smem_capacity_in_bytes("sm_100") + self.num_tmem_alloc_cols = 512 + + # set stages for ab_pipeline and acc_pipeline + self.num_acc_stage = 1 + self.num_ab_stage = 5 + + @cute.jit + def __call__( + self, + a_ptr: cute.Pointer, + b_ptr: cute.Pointer, + sfa_ptr: cute.Pointer, + sfb_ptr: cute.Pointer, + c_ptr: cute.Pointer, + problem_size: tuple, + stream: cuda.CUstream, + epilogue_op: cutlass.Constexpr = lambda x: x, + ): + # setup static attributes before smem/grid/tma computation + self.c_layout = utils.LayoutEnum.ROW_MAJOR + m, n, k, l = problem_size + + self.use_2cta_instrs = False if mma_tiler_mn[0] == 128 else True + + # Setup attributes that depend on gemm inputs + mma_inst_tile_k = 4 + self.mma_tiler = ( + mma_tiler_mn[0], + mma_tiler_mn[1], + mma_inst_shape_k * mma_inst_tile_k, + ) + + self.mma_inst_shape_sfb = ( + mma_tiler_mn[0] // (2 if self.use_2cta_instrs else 1), + mma_tiler_mn[1], + mma_inst_shape_k, + ) + self.mma_tiler_sfb = ( + self.mma_inst_shape_sfb[0], + self.mma_inst_shape_sfb[1], + mma_inst_shape_k * mma_inst_tile_k, + ) + + a_tensor = cute.make_tensor( + a_ptr, + cute.make_layout( + (m, cute.assume(k, 32), l), + stride=(cute.assume(k, 32), 1, cute.assume(m * k, 32)), + ), + ) + b_tensor = cute.make_tensor( + b_ptr, + cute.make_layout( + (n, cute.assume(k, 32), l), + stride=(cute.assume(k, 32), 1, cute.assume(n * k, 32)), + ), + ) + # 256bit aligned. row_major + c_tensor = cute.make_tensor( + c_ptr, + cute.make_layout( + (cute.assume(m, 32), cute.assume(n, 16), l), + stride=(cute.assume(n, 16), 1, cute.assume(m * n, 512)), + ), + ) + + # Setup sfa/sfb tensor by filling A/B tensor to scale factor atom layout + # ((Atom_M, Rest_M),(Atom_K, Rest_K),RestL) + sfa_layout = blockscaled_utils.tile_atom_to_shape_SF( + a_tensor.shape, sf_vec_size + ) + sfa_tensor = cute.make_tensor(sfa_ptr, sfa_layout) + + # ((Atom_N, Rest_N),(Atom_K, Rest_K),RestL) + sfb_layout = blockscaled_utils.tile_atom_to_shape_SF( + b_tensor.shape, sf_vec_size + ) + sfb_tensor = cute.make_tensor(sfb_ptr, sfb_layout) + + mma_op = tcgen05.MmaMXF4NVF4Op( + sf_dtype, + (*mma_tiler_mn, mma_inst_shape_k), + tcgen05.CtaGroup.ONE if not self.use_2cta_instrs else tcgen05.CtaGroup.TWO, + tcgen05.OperandSource.SMEM, + ) + tiled_mma = cute.make_tiled_mma(mma_op) + + # (CTA_Tile_Shape_M, Round_Up(MMA_Tile_Shape_N, 128), MMA_Inst_Shape_K) + # Note sfB don't support share among 2ctas + sfb_mma_op = tcgen05.MmaMXF4NVF4Op( + sf_dtype, + self.mma_inst_shape_sfb, + tcgen05.CtaGroup.ONE, + tcgen05.OperandSource.SMEM, + ) + tiled_mma_sfb = cute.make_tiled_mma(sfb_mma_op) + + self.cta_tile_shape_mnk = ( + self.mma_tiler[0] // (2 if self.use_2cta_instrs else 1), + self.mma_tiler[1], + self.mma_tiler[2], + ) + self.cta_tile_shape_mnk_sfb = ( + self.mma_tiler_sfb[0] // (2 if self.use_2cta_instrs else 1), + self.mma_tiler_sfb[1], + self.mma_tiler_sfb[2], + ) + + self.cluster_layout_vmnk = cute.tiled_divide( + cute.make_layout(cluster_shape_mnk), + (tiled_mma.thr_id.shape,), + ) + self.cluster_layout_sfb_vmnk = cute.tiled_divide( + cute.make_layout(cluster_shape_mnk), + (tiled_mma_sfb.thr_id.shape,), + ) + + # Compute number of multicast CTAs for A/B + self.num_mcast_ctas_a = cute.size(self.cluster_layout_vmnk.shape[2]) + self.num_mcast_ctas_b = cute.size(self.cluster_layout_vmnk.shape[1]) + self.num_mcast_ctas_sfb = cute.size(self.cluster_layout_sfb_vmnk.shape[1]) + self.is_a_mcast = self.num_mcast_ctas_a > 1 + self.is_b_mcast = self.num_mcast_ctas_b > 1 + self.is_sfb_mcast = self.num_mcast_ctas_sfb > 1 + + # Compute A/B/SFA/SFB/C shared memory layout + self.a_smem_layout_staged = sm100_utils.make_smem_layout_a( + tiled_mma, + self.mma_tiler, + ab_dtype, + self.num_ab_stage, + ) + self.b_smem_layout_staged = sm100_utils.make_smem_layout_b( + tiled_mma, + self.mma_tiler, + ab_dtype, + self.num_ab_stage, + ) + self.sfa_smem_layout_staged = blockscaled_utils.make_smem_layout_sfa( + tiled_mma, + self.mma_tiler, + sf_vec_size, + self.num_ab_stage, + ) + self.sfb_smem_layout_staged = blockscaled_utils.make_smem_layout_sfb( + tiled_mma, + self.mma_tiler, + sf_vec_size, + self.num_ab_stage, + ) + + atom_thr_size = cute.size(tiled_mma.thr_id.shape) + + a_op = sm100_utils.cluster_shape_to_tma_atom_A( + cluster_shape_mnk[:2], tiled_mma.thr_id + ) + # TMA load for A + a_smem_layout = cute.slice_(self.a_smem_layout_staged, (None, None, None, 0)) + tma_atom_a, tma_tensor_a = cute.nvgpu.make_tiled_tma_atom_A( + a_op, + a_tensor, + a_smem_layout, + self.mma_tiler, + tiled_mma, + self.cluster_layout_vmnk.shape, + ) + # TMA load for B + b_op = sm100_utils.cluster_shape_to_tma_atom_B( + cluster_shape_mnk[:2], tiled_mma.thr_id + ) + b_smem_layout = cute.slice_(self.b_smem_layout_staged, (None, None, None, 0)) + tma_atom_b, tma_tensor_b = cute.nvgpu.make_tiled_tma_atom_B( + b_op, + b_tensor, + b_smem_layout, + self.mma_tiler, + tiled_mma, + self.cluster_layout_vmnk.shape, + ) + + # TMA load for SFA + sfa_op = sm100_utils.cluster_shape_to_tma_atom_A( + cluster_shape_mnk[:2], tiled_mma.thr_id + ) + sfa_smem_layout = cute.slice_( + self.sfa_smem_layout_staged, (None, None, None, 0) + ) + tma_atom_sfa, tma_tensor_sfa = cute.nvgpu.make_tiled_tma_atom_A( + sfa_op, + sfa_tensor, + sfa_smem_layout, + self.mma_tiler, + tiled_mma, + self.cluster_layout_vmnk.shape, + internal_type=cutlass.Int16, + ) + + # TMA load for SFB + sfb_op = sm100_utils.cluster_shape_to_tma_atom_SFB( + cluster_shape_mnk[:2], tiled_mma.thr_id + ) + sfb_smem_layout = cute.slice_( + self.sfb_smem_layout_staged, (None, None, None, 0) + ) + tma_atom_sfb, tma_tensor_sfb = cute.nvgpu.make_tiled_tma_atom_B( + sfb_op, + sfb_tensor, + sfb_smem_layout, + self.mma_tiler_sfb, + tiled_mma_sfb, + self.cluster_layout_sfb_vmnk.shape, + internal_type=cutlass.Int16, + ) + + # Compute TMA load bytes + a_copy_size = cute.size_in_bytes(ab_dtype, a_smem_layout) + b_copy_size = cute.size_in_bytes(ab_dtype, b_smem_layout) + sfa_copy_size = cute.size_in_bytes(sf_dtype, sfa_smem_layout) + sfb_copy_size = cute.size_in_bytes(sf_dtype, sfb_smem_layout) + self.num_tma_load_bytes = ( + a_copy_size + b_copy_size + sfa_copy_size + sfb_copy_size + ) * atom_thr_size + + # Compute grid size + grid = cute.round_up( + cute.ceil_div( + (c_tensor.layout.shape), + (self.cta_tile_shape_mnk[0], self.cta_tile_shape_mnk[1], 1), + ), + cluster_shape_mnk, + ) + + # Launch the kernel + self.kernel( + tiled_mma, + tiled_mma_sfb, + tma_atom_a, + tma_tensor_a, + tma_atom_b, + tma_tensor_b, + tma_atom_sfa, + tma_tensor_sfa, + tma_atom_sfb, + tma_tensor_sfb, + c_tensor, + self.a_smem_layout_staged, + self.b_smem_layout_staged, + self.sfa_smem_layout_staged, + self.sfb_smem_layout_staged, + self.cluster_layout_vmnk, + self.cluster_layout_sfb_vmnk, + epilogue_op, + ).launch( + grid=grid, + block=[self.threads_per_cta, 1, 1], + cluster=cluster_shape_mnk, + stream=stream, + ) + return + + # GPU device kernel + @cute.kernel + def kernel( + self, + tiled_mma: cute.TiledMma, + tiled_mma_sfb: cute.TiledMma, + tma_atom_a: cute.CopyAtom, + mA_mkl: cute.Tensor, + tma_atom_b: cute.CopyAtom, + mB_nkl: cute.Tensor, + tma_atom_sfa: cute.CopyAtom, + mSFA_mkl: cute.Tensor, + tma_atom_sfb: cute.CopyAtom, + mSFB_nkl: cute.Tensor, + mC_mnl: cute.Tensor, + a_smem_layout_staged: cute.ComposedLayout, + b_smem_layout_staged: cute.ComposedLayout, + sfa_smem_layout_staged: cute.Layout, + sfb_smem_layout_staged: cute.Layout, + cta_layout_vmnk: cute.Layout, + cta_layout_sfb_vmnk: cute.Layout, + epilogue_op: cutlass.Constexpr, + ): + """ + GPU device kernel performing the batched GEMM computation. + """ + warp_idx = cute.arch.warp_idx() + warp_idx = cute.arch.make_warp_uniform(warp_idx) + tidx, _, _ = cute.arch.thread_idx() + + # + # Setup cta/thread coordinates + # + # Coords inside cluster + bidx, bidy, bidz = cute.arch.block_idx() + cta_rank_in_cluster = cute.arch.block_idx_in_cluster() + cta_in_cluster_coord_vmnk = cta_layout_vmnk.get_flat_coord(cta_rank_in_cluster) + cta_in_cluster_coord_sfb_vmnk = cta_layout_sfb_vmnk.get_flat_coord( + cta_rank_in_cluster + ) + + # Coords outside cluster + mma_tile_coord_vmnk = ( + bidx % cute.size(cta_layout_vmnk, mode=[0]), + bidx // cute.size(cta_layout_vmnk, mode=[0]), + bidy, + bidz, + ) + mma_tile_coord_mnl = mma_tile_coord_vmnk[1:] + is_leader_cta = mma_tile_coord_vmnk[0] == 0 + + # + # Define shared storage for kernel + # + @cute.struct + class SharedStorage: + ab_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] + acc_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] + tmem_dealloc_mbar_ptr: cutlass.Int64 + tmem_holding_buf: cutlass.Int32 + + smem = utils.SmemAllocator() + storage = smem.allocate(SharedStorage) + # (MMA, MMA_M, MMA_K, STAGE) + sA = smem.allocate_tensor( + element_type=ab_dtype, + layout=a_smem_layout_staged.outer, + byte_alignment=128, + swizzle=a_smem_layout_staged.inner, + ) + # (MMA, MMA_N, MMA_K, STAGE) + sB = smem.allocate_tensor( + element_type=ab_dtype, + layout=b_smem_layout_staged.outer, + byte_alignment=128, + swizzle=b_smem_layout_staged.inner, + ) + # (MMA, MMA_M, MMA_K, STAGE) + sSFA = smem.allocate_tensor( + element_type=sf_dtype, + layout=sfa_smem_layout_staged, + byte_alignment=128, + ) + # (MMA, MMA_N, MMA_K, STAGE) + sSFB = smem.allocate_tensor( + element_type=sf_dtype, + layout=sfb_smem_layout_staged, + byte_alignment=128, + ) + + # + # Compute multicast mask for A/B/SFA/SFB buffer full + # + a_full_mcast_mask = None + b_full_mcast_mask = None + sfa_full_mcast_mask = None + sfb_full_mcast_mask = None + if cutlass.const_expr( + self.is_a_mcast or self.is_b_mcast or self.use_2cta_instrs + ): + a_full_mcast_mask = cpasync.create_tma_multicast_mask( + cta_layout_vmnk, cta_in_cluster_coord_vmnk, mcast_mode=2 + ) + b_full_mcast_mask = cpasync.create_tma_multicast_mask( + cta_layout_vmnk, cta_in_cluster_coord_vmnk, mcast_mode=1 + ) + sfa_full_mcast_mask = cpasync.create_tma_multicast_mask( + cta_layout_vmnk, cta_in_cluster_coord_vmnk, mcast_mode=2 + ) + sfb_full_mcast_mask = cpasync.create_tma_multicast_mask( + cta_layout_sfb_vmnk, cta_in_cluster_coord_sfb_vmnk, mcast_mode=1 + ) + + # + # Initialize mainloop ab_pipeline, acc_pipeline and their states + # + num_tma_producer = self.num_mcast_ctas_a + self.num_mcast_ctas_b - 1 + ab_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) + ab_pipeline_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, num_tma_producer + ) + ab_producer, ab_consumer = pipeline.PipelineTmaUmma.create( + barrier_storage=storage.ab_mbar_ptr.data_ptr(), + num_stages=self.num_ab_stage, + producer_group=ab_pipeline_producer_group, + consumer_group=ab_pipeline_consumer_group, + tx_count=self.num_tma_load_bytes, + cta_layout_vmnk=cta_layout_vmnk, + ).make_participants() + acc_producer, acc_consumer = pipeline.PipelineUmmaAsync.create( + barrier_storage=storage.acc_mbar_ptr.data_ptr(), + num_stages=self.num_acc_stage, + producer_group=ab_pipeline_producer_group, + consumer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, + self.threads_per_cta * (2 if self.use_2cta_instrs else 1), + ), + cta_layout_vmnk=cta_layout_vmnk, + ).make_participants() + + # + # Local_tile partition global tensors + # + # (bM, bK, RestM, RestK, RestL) + gA_mkl = cute.local_tile( + mA_mkl, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None) + ) + # (bN, bK, RestN, RestK, RestL) + gB_nkl = cute.local_tile( + mB_nkl, cute.slice_(self.mma_tiler, (0, None, None)), (None, None, None) + ) + gSFA_mkl = cute.local_tile( + mSFA_mkl, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None) + ) + gSFB_nkl = cute.local_tile( + mSFB_nkl, cute.slice_(self.mma_tiler, (0, None, None)), (None, None, None) + ) + # (bM, bN, RestM, RestN, RestL) + gC_mnl = cute.local_tile( + mC_mnl, cute.slice_(self.mma_tiler, (None, None, 0)), (None, None, None) + ) + k_tile_cnt = cute.size(gA_mkl, mode=[3]) + + # + # Partition global tensor for TiledMMA_A/B/SFA/SFB/C + # + # (MMA, MMA_M, MMA_K, RestK) + thr_mma = tiled_mma.get_slice(mma_tile_coord_vmnk[0]) + thr_mma_sfb = tiled_mma_sfb.get_slice(mma_tile_coord_vmnk[0]) + # (MMA, MMA_M, MMA_K, RestM, RestK, RestL) + tCgA = thr_mma.partition_A(gA_mkl) + # (MMA, MMA_N, MMA_K, RestN, RestK, RestL) + tCgB = thr_mma.partition_B(gB_nkl) + # (MMA, MMA_M, MMA_K, RestM, RestK, RestL) + tCgSFA = thr_mma.partition_A(gSFA_mkl) + # (MMA, MMA_N, MMA_K, RestN, RestK, RestL) + # tCgSFB = thr_mma.partition_B(gSFB_nkl) + tCgSFB = thr_mma_sfb.partition_B(gSFB_nkl) + # (MMA, MMA_M, MMA_N, RestM, RestN, RestL) + tCgC = thr_mma.partition_C(gC_mnl) + + # + # Partition global/shared tensor for TMA load A/B/SFA/SFB + # + # TMA load A partition_S/D + # ((atom_v, rest_v), STAGE) + # ((atom_v, rest_v), RestM, RestK, RestL) + tAsA, tAgA = cpasync.tma_partition( + tma_atom_a, + # 0, + # cute.make_layout(1), + cta_in_cluster_coord_vmnk[2], + cute.make_layout(cute.size(cta_layout_vmnk, mode=[2])), + cute.group_modes(sA, 0, 3), + cute.group_modes(tCgA, 0, 3), + ) + # TMA load B partition_S/D + # ((atom_v, rest_v), STAGE) + # ((atom_v, rest_v), RestN, RestK, RestL) + tBsB, tBgB = cpasync.tma_partition( + tma_atom_b, + # 0, + # cute.make_layout(1), + cta_in_cluster_coord_vmnk[1], + cute.make_layout(cute.size(cta_layout_vmnk, mode=[1])), + cute.group_modes(sB, 0, 3), + cute.group_modes(tCgB, 0, 3), + ) + + # TMA load SFA partition_S/D + # ((atom_v, rest_v), STAGE) + # ((atom_v, rest_v), RestM, RestK, RestL) + tAsSFA, tAgSFA = cpasync.tma_partition( + tma_atom_sfa, + # 0, + # cute.make_layout(1), + cta_in_cluster_coord_vmnk[2], + cute.make_layout(cute.size(cta_layout_vmnk, mode=[2])), + cute.group_modes(sSFA, 0, 3), + cute.group_modes(tCgSFA, 0, 3), + ) + tAsSFA = cute.filter_zeros(tAsSFA) + tAgSFA = cute.filter_zeros(tAgSFA) + + # TMA load SFB partition_S/D + # ((atom_v, rest_v), STAGE) + # ((atom_v, rest_v), RestN, RestK, RestL) + sfb_cta_layout = cute.make_layout( + cute.slice_(cta_layout_sfb_vmnk, (0, None, 0, 0)).shape + ) + tBsSFB, tBgSFB = cpasync.tma_partition( + tma_atom_sfb, + cta_in_cluster_coord_sfb_vmnk[1], + sfb_cta_layout, + cute.group_modes(sSFB, 0, 3), + cute.group_modes(tCgSFB, 0, 3), + ) + tBsSFB = cute.filter_zeros(tBsSFB) + tBgSFB = cute.filter_zeros(tBgSFB) + + # + # Partition shared/tensor memory tensor for TiledMMA_A/B/C + # + # (MMA, MMA_M, MMA_K, STAGE) + tCrA = tiled_mma.make_fragment_A(sA) + # (MMA, MMA_N, MMA_K, STAGE) + tCrB = tiled_mma.make_fragment_B(sB) + # (MMA, MMA_M, MMA_N) + acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2]) + # (MMA, MMA_M, MMA_N) + tCtAcc_fake = tiled_mma.make_fragment_C(acc_shape) + + # + # Alloc tensor memory buffer + # + tmem_alloc_barrier = pipeline.NamedBarrier( + barrier_id=1, + num_threads=self.threads_per_cta, + ) + tmem = utils.TmemAllocator( + storage.tmem_holding_buf, + barrier_for_retrieve=tmem_alloc_barrier, + is_two_cta=cute.size(cta_layout_vmnk, mode=[0]) > 1, + two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar_ptr, + ) + tmem.allocate(self.num_tmem_alloc_cols) + tmem.wait_for_alloc() + acc_tmem_ptr = tmem.retrieve_ptr(cutlass.Float32) + tCtAcc = cute.make_tensor(acc_tmem_ptr, tCtAcc_fake.layout) + + # + # Make SFA/SFB tmem tensor + # + # Get SFA tmem ptr + sfa_tmem_ptr = cute.recast_ptr( + acc_tmem_ptr + tcgen05.find_tmem_tensor_col_offset(tCtAcc), + dtype=sf_dtype, + ) + # (MMA, MMA_M, MMA_K) + tCtSFA_layout = blockscaled_utils.make_tmem_layout_sfa( + tiled_mma, + self.mma_tiler, + sf_vec_size, + cute.slice_(sfa_smem_layout_staged, (None, None, None, 0)), + ) + tCtSFA = cute.make_tensor(sfa_tmem_ptr, tCtSFA_layout) + # Get SFB tmem ptr + sfb_tmem_ptr = cute.recast_ptr( + acc_tmem_ptr + + tcgen05.find_tmem_tensor_col_offset(tCtAcc) + + tcgen05.find_tmem_tensor_col_offset(tCtSFA), + dtype=sf_dtype, + ) + # (MMA, MMA_N, MMA_K) + tCtSFB_layout = blockscaled_utils.make_tmem_layout_sfb( + tiled_mma, + self.mma_tiler, + sf_vec_size, + cute.slice_(sfb_smem_layout_staged, (None, None, None, 0)), + ) + tCtSFB = cute.make_tensor(sfb_tmem_ptr, tCtSFB_layout) + + # + # Partition for S2T copy of SFA/SFB + # + # Make S2T CopyAtom + copy_atom_s2t = cute.make_copy_atom( + tcgen05.Cp4x32x128bOp( + tcgen05.CtaGroup.ONE + if not self.use_2cta_instrs + else tcgen05.CtaGroup.TWO + ), + sf_dtype, + ) + # (MMA, MMA_MN, MMA_K, STAGE) + tCsSFA_compact = cute.filter_zeros(sSFA) + # (MMA, MMA_MN, MMA_K) + tCtSFA_compact = cute.filter_zeros(tCtSFA) + tiled_copy_s2t_sfa = tcgen05.make_s2t_copy(copy_atom_s2t, tCtSFA_compact) + thr_copy_s2t_sfa = tiled_copy_s2t_sfa.get_slice(0) + # ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K, STAGE) + tCsSFA_compact_s2t_ = thr_copy_s2t_sfa.partition_S(tCsSFA_compact) + # ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K, STAGE) + tCsSFA_compact_s2t = tcgen05.get_s2t_smem_desc_tensor( + tiled_copy_s2t_sfa, tCsSFA_compact_s2t_ + ) + # ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K) + tCtSFA_compact_s2t = thr_copy_s2t_sfa.partition_D(tCtSFA_compact) + + # (MMA, MMA_MN, MMA_K, STAGE) + tCsSFB_compact = cute.filter_zeros(sSFB) + # (MMA, MMA_MN, MMA_K) + tCtSFB_compact = cute.filter_zeros(tCtSFB) + tiled_copy_s2t_sfb = tcgen05.make_s2t_copy(copy_atom_s2t, tCtSFB_compact) + thr_copy_s2t_sfb = tiled_copy_s2t_sfb.get_slice(0) + # ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K, STAGE) + tCsSFB_compact_s2t_ = thr_copy_s2t_sfb.partition_S(tCsSFB_compact) + # ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K, STAGE) + tCsSFB_compact_s2t = tcgen05.get_s2t_smem_desc_tensor( + tiled_copy_s2t_sfb, tCsSFB_compact_s2t_ + ) + # ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K) + tCtSFB_compact_s2t = thr_copy_s2t_sfb.partition_D(tCtSFB_compact) + + # + # Slice to per mma tile index + # + # ((atom_v, rest_v), RestK) + tAgA = tAgA[(None, mma_tile_coord_mnl[0], None, mma_tile_coord_mnl[2])] + # ((atom_v, rest_v), RestK) + tBgB = tBgB[(None, mma_tile_coord_mnl[1], None, mma_tile_coord_mnl[2])] + # ((atom_v, rest_v), RestK) + tAgSFA = tAgSFA[(None, mma_tile_coord_mnl[0], None, mma_tile_coord_mnl[2])] + # ((atom_v, rest_v), RestK) + tBgSFB = tBgSFB[(None, mma_tile_coord_mnl[1], None, mma_tile_coord_mnl[2])] + + # + # Execute Data copy and Math computation in the k_tile loop + # + if warp_idx == 0: + # Wait for accumulator buffer empty + if is_leader_cta: + acc_producer.acquire_and_advance() + + # Set ACCUMULATE field to False for the first k_tile iteration + tiled_mma.set(tcgen05.Field.ACCUMULATE, False) + # Execute k_tile loop + for k_tile in cutlass.range( + k_tile_cnt, prefetch_stages=self.num_ab_stage - 2 + ): + # Wait for AB buffer empty + ab_empty = ab_producer.acquire_and_advance() + + # TMA load A/B/SFA/SFB + cute.copy( + tma_atom_a, + tAgA[(None, ab_empty.count)], + tAsA[(None, ab_empty.index)], + tma_bar_ptr=ab_empty.barrier, + mcast_mask=a_full_mcast_mask, + ) + cute.copy( + tma_atom_b, + tBgB[(None, ab_empty.count)], + tBsB[(None, ab_empty.index)], + tma_bar_ptr=ab_empty.barrier, + mcast_mask=b_full_mcast_mask, + ) + cute.copy( + tma_atom_sfa, + tAgSFA[(None, ab_empty.count)], + tAsSFA[(None, ab_empty.index)], + tma_bar_ptr=ab_empty.barrier, + mcast_mask=sfa_full_mcast_mask, + ) + cute.copy( + tma_atom_sfb, + tBgSFB[(None, ab_empty.count)], + tBsSFB[(None, ab_empty.index)], + tma_bar_ptr=ab_empty.barrier, + mcast_mask=sfb_full_mcast_mask, + ) + + if is_leader_cta: + # Wait for AB buffer full + ab_full = ab_consumer.wait_and_advance() + + # Copy SFA/SFB to tmem + s2t_stage_coord = (None, None, None, None, ab_full.index) + tCsSFA_compact_s2t_staged = tCsSFA_compact_s2t[s2t_stage_coord] + tCsSFB_compact_s2t_staged = tCsSFB_compact_s2t[s2t_stage_coord] + cute.copy( + tiled_copy_s2t_sfa, + tCsSFA_compact_s2t_staged, + tCtSFA_compact_s2t, + ) + cute.copy( + tiled_copy_s2t_sfb, + tCsSFB_compact_s2t_staged, + tCtSFB_compact_s2t, + ) + + # tCtAcc += tCrA * tCrSFA * tCrB * tCrSFB + num_kblocks = cute.size(tCrA, mode=[2]) + for kblock_idx in cutlass.range(num_kblocks, unroll_full=True): + kblock_coord = ( + None, + None, + kblock_idx, + ab_full.index, + ) + + # Set SFA/SFB tensor to tiled_mma + sf_kblock_coord = (None, None, kblock_idx) + tiled_mma.set( + tcgen05.Field.SFA, + tCtSFA[sf_kblock_coord].iterator, + ) + tiled_mma.set( + tcgen05.Field.SFB, + tCtSFB[sf_kblock_coord].iterator, + ) + + cute.gemm( + tiled_mma, + tCtAcc, + tCrA[kblock_coord], + tCrB[kblock_coord], + tCtAcc, + ) + # Enable accumulate on tCtAcc after first kblock + tiled_mma.set(tcgen05.Field.ACCUMULATE, True) + + # Async arrive AB buffer empty + ab_full.release() + if is_leader_cta: + acc_producer.commit() + + # + # Epilogue + # Partition for epilogue + # + # x32 or x128 all is ok. + op = tcgen05.Ld32x32bOp(tcgen05.Repetition.x128, tcgen05.Pack.NONE) + copy_atom_t2r = cute.make_copy_atom(op, cutlass.Float32) + tiled_copy_t2r = tcgen05.make_tmem_copy(copy_atom_t2r, tCtAcc) + thr_copy_t2r = tiled_copy_t2r.get_slice(tidx) + # (T2R_M, T2R_N, EPI_M, EPI_M) + tTR_tAcc = thr_copy_t2r.partition_S(tCtAcc) + # (T2R_M, T2R_N, EPI_M, EPI_N, RestM, RestN, RestL) + tTR_gC = thr_copy_t2r.partition_D(tCgC) + # (T2R_M, T2R_N, EPI_M, EPI_N） + tTR_rAcc = cute.make_rmem_tensor( + tTR_gC[None, None, None, None, 0, 0, 0].shape, cutlass.Float32 + ) + # (T2R_M, T2R_N, EPI_M, EPI_N） + tTR_rC = cute.make_rmem_tensor( + tTR_gC[None, None, None, None, 0, 0, 0].shape, c_dtype + ) + # STG Atom + simt_atom = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), c_dtype) + tTR_gC = tTR_gC[(None, None, None, None, *mma_tile_coord_mnl)] + + # Wait for accumulator buffer full + acc_full = acc_consumer.wait_and_advance() + + # Copy accumulator to register + cute.copy(tiled_copy_t2r, tTR_tAcc, tTR_rAcc) + acc_vec = epilogue_op(tTR_rAcc.load().to(c_dtype)) + tTR_rC.store(acc_vec) + # Store C to global memory + cute.copy(simt_atom, tTR_rC, tTR_gC) + + acc_full.release() + + # Ensure used buffers are properly synchronized before producer exit. + # This could avoid the invalid dsmem access due to early leading CTA exit. + if warp_idx == 0: + ab_producer.tail() + if is_leader_cta: + acc_producer.tail() + + # Deallocate TMEM + cute.arch.barrier() + tmem.free(acc_tmem_ptr) + + return + + +def run_nvfp4_gemm( + mnkl: Tuple[int, int, int, int], + tolerance: float, + warmup_iterations: int = 10, + iterations: int = 100, + use_cold_l2: bool = True, + do_benchmark: bool = False, +): + run( + gemm_class=Sm100BlockScaledDenseGemmKernel, + ab_dtype=ab_dtype, + sf_dtype=sf_dtype, + c_dtype=c_dtype, + sf_vec_size=sf_vec_size, + mma_tiler_mn=mma_tiler_mn, + cluster_shape_mnk=cluster_shape_mnk, + mnkl=mnkl, + tolerance=tolerance, + do_benchmark=do_benchmark, + warmup_iterations=warmup_iterations, + iterations=iterations, + use_cold_l2=use_cold_l2, + ) + + +if __name__ == "__main__": + parser = create_parser() + args = parser.parse_args() + if len(args.mnkl) != 4: + parser.error("--mnkl must contain exactly 4 values") + + m, n, k, _ = args.mnkl + if m % mma_tiler_mn[0] != 0: + parser.error("M must be multiples of mma_tiler_mn[0] (got m={})".format(m)) + if n % mma_tiler_mn[1] != 0: + parser.error("N must be multiples of mma_tiler_mn[1] (got n={})".format(n)) + if k % 256 != 0: + parser.error("k must be a multiple of 256 (got k={})".format(k)) + + run_nvfp4_gemm( + args.mnkl, + args.tolerance, + do_benchmark=args.do_benchmark, + ) + print("PASS") diff --git a/examples/python/CuTeDSL/blackwell/tutorial_gemm/utils.py b/examples/python/CuTeDSL/blackwell/tutorial_gemm/utils.py new file mode 100644 index 00000000..f9c96b20 --- /dev/null +++ b/examples/python/CuTeDSL/blackwell/tutorial_gemm/utils.py @@ -0,0 +1,366 @@ +# 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. + +import argparse +from typing import Tuple + +import torch + +import cutlass +import cutlass.cute as cute +import cutlass.torch as cutlass_torch +from cutlass.cute.runtime import make_ptr + + +def parse_comma_separated_ints(s: str) -> Tuple[int, ...]: + try: + return tuple(int(x.strip()) for x in s.split(",")) + except ValueError: + raise argparse.ArgumentTypeError( + "Invalid format. Expected comma-separated integers." + ) + + +def create_parser(): + parser = argparse.ArgumentParser( + description="Example of Sm100 Dense BlockScaled GEMM." + ) + parser.add_argument( + "--mnkl", + type=parse_comma_separated_ints, + default=(8192, 8192, 8192, 8), + help="mnkl dimensions (comma-separated)", + ) + parser.add_argument( + "--tolerance", type=float, default=1e-01, help="Tolerance for validation" + ) + parser.add_argument( + "--do_benchmark", action="store_true", default=False, help="Do benchmark test" + ) + return parser + + +def ceil_div(a, b): + return (a + b - 1) // b + + +# Helper function to create scale factor tensor SFA/SFB +# for 1x16 block scaled wise use case and follow the layout requirement +# defined in https://docs.nvidia.com/cuda/cublas/index.html?highlight=fp4#d-block-scaling-factors-layout +@cute.jit +def cvt_sf_MKL_to_M32x4xrm_K4xrk_L( + sf_ref_ptr: cute.Pointer, + sf_mma_ptr: cute.Pointer, + mn: int, + sf_k: int, + l: int, + mma_shape: tuple, +): + mma_permute_order = (3, 4, 1, 5, 2, 0) + permuted_shape = tuple(mma_shape[i] for i in mma_permute_order) + cute_layout = cute.make_ordered_layout(permuted_shape, order=(2, 1, 4, 0, 3, 5)) + + sf_ref_tensor = cute.make_tensor( + sf_ref_ptr, cute.make_layout((mn, sf_k, l), stride=(sf_k, 1, mn * sf_k)) + ) + sf_mma_tensor = cute.make_tensor(sf_mma_ptr, cute_layout) + + sf_mma_tensor = cute.group_modes(sf_mma_tensor, 0, 3) + sf_mma_tensor = cute.group_modes(sf_mma_tensor, 1, 3) + for i in cutlass.range(cute.size(sf_ref_tensor)): + mkl_coord = sf_ref_tensor.layout.get_hier_coord(i) + sf_mma_tensor[mkl_coord] = sf_ref_tensor[mkl_coord] + pass + + +def to_blocked(input_matrix): + rows, cols = input_matrix.shape + + # Please ensure rows and cols are multiples of 128 and 4 respectively + n_row_blocks = ceil_div(rows, 128) + n_col_blocks = ceil_div(cols, 4) + + padded = input_matrix + blocks = padded.view(n_row_blocks, 128, n_col_blocks, 4).permute(0, 2, 1, 3) + rearranged = blocks.reshape(-1, 4, 32, 4).transpose(1, 2).reshape(-1, 32, 16) + + return rearranged.flatten() + + +def run( + gemm_class, + ab_dtype, + sf_dtype, + c_dtype, + sf_vec_size, + mma_tiler_mn, + cluster_shape_mnk, + mnkl: Tuple[int, int, int, int], + tolerance: float, + warmup_iterations: int = 10, + iterations: int = 100, + use_cold_l2: bool = True, + do_benchmark: bool = False, +): + """ + Prepare A/B/SFA/SFB/C tensors, launch GPU kernel, and reference checking. + """ + print("=" * 60) + print("Launching Blackwell Dense BlockScaled GEMM Test") + print("-" * 60) + print(f"Input dimensions (m, n, k, l): {mnkl}") + print(f" m (rows): {mnkl[0]}") + print(f" n (cols): {mnkl[1]}") + print(f" k (inner): {mnkl[2]}") + print(f" l (batch): {mnkl[3]}") + print(f"Data Types & Precision:") + print(f" Input matrices (A, B): {ab_dtype}") + print(f" Scale factors (SFA, SFB): {sf_dtype}") + print(f" Output matrix (C): {c_dtype}") + print(f" Scale factor vector size: {sf_vec_size}") + print("Tile and cluster configuration:") + print(f" MMA tiler (M, N, K): {mma_tiler_mn}") + print(f" Cluster shape (M, N, K): {cluster_shape_mnk}") + print(f"Validation tolerance: {tolerance}") + print(f"Do benchmark: {do_benchmark}") + print("=" * 60) + + # Unpack parameters + m, n, k, l = mnkl + + if not torch.cuda.is_available(): + raise RuntimeError("GPU is required to run this example!") + + torch.manual_seed(1111) + + # Create tensor A/B/C + a_ref = torch.randint( + 0, 2, (l, m, k // 2), dtype=torch.uint8, device="cuda" + ).permute(1, 2, 0) + b_ref = torch.randint( + 0, 2, (l, n, k // 2), dtype=torch.uint8, device="cuda" + ).permute(1, 2, 0) + # a_ref = torch.ones((l, m, k // 2), dtype=torch.uint8, device="cuda").permute(1, 2, 0) + # b_ref = torch.ones((l, n, k // 2), dtype=torch.uint8, device="cuda").permute(1, 2, 0) + a_ref_f4 = a_ref.view(torch.float4_e2m1fn_x2) + b_ref_f4 = b_ref.view(torch.float4_e2m1fn_x2) + + c_tensor = torch.randn((l, m, n), dtype=torch.float16, device="cuda").permute( + 1, 2, 0 + ) + + # Create a torch tensor for scale factor tensor of A and B + def create_ref_scale_factor_tensor(l, mn, sf_k): + """ + Create the reference scale factor tensor on CPU. + Returns the reshaped/pruned tensor ready for ref computation and its original permuted form. + """ + ref_shape = (l, mn, sf_k) + ref_permute_order = (1, 2, 0) + ref_f8_random_int = torch.randint(1, 3, ref_shape, dtype=torch.int8) + ref_f8_torch_tensor_cpu = ref_f8_random_int.to(dtype=torch.float8_e4m3fn) + # permute to match ref_permute_order + ref_f8_torch_tensor_cpu_permuted = ref_f8_torch_tensor_cpu.permute( + *ref_permute_order + ) + return ref_f8_torch_tensor_cpu_permuted + + # Copy the reference scale factor tensor to the CUTE-format scale factor tensor + def create_cute_scale_factor_tensor(l, mn, sf_k, ref_f8_torch_tensor_cpu_permuted): + """ + Create the CUTE-format scale factor tensor on CUDA based on the reference tensor. + """ + atom_m = (32, 4) + atom_k = 4 + mma_shape = ( + l, # batch size + ceil_div(mn, atom_m[0] * atom_m[1]), + ceil_div(sf_k, atom_k), + atom_m[0], + atom_m[1], + atom_k, + ) + mma_permute_order = (3, 4, 1, 5, 2, 0) + + # Generate a random int8 tensor, then convert to float8_e4m3fn + rand_int_tensor = torch.randint(0, 2, mma_shape, dtype=torch.int8) + cute_f8_torch_tensor_cpu = rand_int_tensor.to(dtype=torch.float8_e4m3fn) + # Permute according to mma_permute_order + cute_f8_torch_tensor_cpu = cute_f8_torch_tensor_cpu.permute(*mma_permute_order) + + # Call the helper function to do layout conversion + cvt_sf_MKL_to_M32x4xrm_K4xrk_L( + make_ptr( + cutlass.Float8E4M3FN, + ref_f8_torch_tensor_cpu_permuted.data_ptr(), + cute.AddressSpace.gmem, + assumed_align=32, + ), + make_ptr( + cutlass.Float8E4M3FN, + cute_f8_torch_tensor_cpu.data_ptr(), + cute.AddressSpace.gmem, + assumed_align=32, + ), + mn, + sf_k, + l, + mma_shape, + ) + return cute_f8_torch_tensor_cpu.cuda() + + sf_k = ceil_div(k, sf_vec_size) + sfa_ref = create_ref_scale_factor_tensor(l, m, sf_k) + sfb_ref = create_ref_scale_factor_tensor(l, n, sf_k) + # sfa_ref.fill_(1) + # sfb_ref.fill_(1) + sfa_tensor = create_cute_scale_factor_tensor(l, m, sf_k, sfa_ref) + sfb_tensor = create_cute_scale_factor_tensor(l, n, sf_k, sfb_ref) + + # Configure gemm kernel + gemm = gemm_class() + # Initialize Stream + current_stream = cutlass_torch.default_stream() + a_ptr = make_ptr( + ab_dtype, a_ref_f4.data_ptr(), cute.AddressSpace.gmem, assumed_align=16 + ) + b_ptr = make_ptr( + ab_dtype, b_ref_f4.data_ptr(), cute.AddressSpace.gmem, assumed_align=16 + ) + c_ptr = make_ptr( + c_dtype, c_tensor.data_ptr(), cute.AddressSpace.gmem, assumed_align=32 + ) + sfa_ptr = make_ptr( + sf_dtype, sfa_tensor.data_ptr(), cute.AddressSpace.gmem, assumed_align=32 + ) + sfb_ptr = make_ptr( + sf_dtype, sfb_tensor.data_ptr(), cute.AddressSpace.gmem, assumed_align=32 + ) + # Compile gemm kernel + compiled_gemm = cute.compile( + gemm, + a_ptr, + b_ptr, + sfa_ptr, + sfb_ptr, + c_ptr, + (m, n, k, l), + current_stream, + ) + # Launch GPU kernel + compiled_gemm(a_ptr, b_ptr, sfa_ptr, sfb_ptr, c_ptr, (m, n, k, l), current_stream) + # For batch l, do (m, k, l) @ (n, k, l).T along k for each batch. + # Result: (m, n, l) + # Allocate ref as (l, m, n) with n-contiguous layout, then permute to (m, n, l) + ref = torch.empty( + (l, m, n), + dtype=torch.float16, + device="cuda", + ).permute(1, 2, 0) + for l_idx in range(l): + # Convert the scale factor tensor to blocked format + scale_a = to_blocked(sfa_ref[:, :, l_idx]) + scale_b = to_blocked(sfb_ref[:, :, l_idx]) + # (m, k) @ (n, k).T -> (m, n) + res = torch._scaled_mm( + a_ref_f4[:, :, l_idx], + b_ref_f4[:, :, l_idx].transpose(0, 1), + scale_a.cuda(), + scale_b.cuda(), + bias=None, + out_dtype=torch.float16, + ) + ref[:, :, l_idx] = res + torch.testing.assert_close(c_tensor, ref, atol=tolerance, rtol=1e-02) + + if do_benchmark: + + def generate_tensors(): + a_ptr = make_ptr( + ab_dtype, a_ref_f4.data_ptr(), cute.AddressSpace.gmem, assumed_align=16 + ) + b_ptr = make_ptr( + ab_dtype, b_ref_f4.data_ptr(), cute.AddressSpace.gmem, assumed_align=16 + ) + c_ptr = make_ptr( + c_dtype, c_tensor.data_ptr(), cute.AddressSpace.gmem, assumed_align=32 + ) + sfa_ptr = make_ptr( + sf_dtype, + sfa_tensor.data_ptr(), + cute.AddressSpace.gmem, + assumed_align=32, + ) + sfb_ptr = make_ptr( + sf_dtype, + sfb_tensor.data_ptr(), + cute.AddressSpace.gmem, + assumed_align=32, + ) + args = cute.testing.JitArguments( + a_ptr, b_ptr, sfa_ptr, sfb_ptr, c_ptr, (m, n, k, l), current_stream + ) + args.add_to_scope([a_ref_f4, b_ref_f4, sfa_tensor, sfb_tensor, c_tensor]) + return args + + workspace_count = 1 + if use_cold_l2: + one_workspace_bytes = ( + a_ref_f4.numel() * a_ref_f4.element_size() + + b_ref_f4.numel() * b_ref_f4.element_size() + + sfa_tensor.numel() * sfa_tensor.element_size() + + sfb_tensor.numel() * sfb_tensor.element_size() + + c_tensor.numel() * c_tensor.element_size() + ) + workspace_count = cute.testing.get_workspace_count( + one_workspace_bytes, warmup_iterations, iterations + ) + + # Return execution time in microseconds + time = cute.testing.benchmark( + compiled_gemm, + workspace_generator=generate_tensors, + workspace_count=workspace_count, + stream=current_stream, + warmup_iterations=warmup_iterations, + iterations=iterations, + ) + + print(f"Execution time: {time} us") + peta_flops = (4 * m * n * k * l) / (time * 1e-6) / 1e9 / 1000000 + print(f"FLOPS: {peta_flops} PFLOPS") + bytes_transfer = ( + 2 * m * k / 2 * l * a_ref_f4.element_size() + + 2 * n * k / 2 * l * b_ref_f4.element_size() + + 2 * m * n * l * c_tensor.element_size() + + 2 * m * sf_k * l * sfa_tensor.element_size() + + 2 * n * sf_k * l * sfb_tensor.element_size() + ) + print(f"Bytes: {bytes_transfer} Bytes") + bandwidth = bytes_transfer / time * 1e-3 + print(f"BW: {bandwidth} GB/s") diff --git a/examples/python/CuTeDSL/blackwell_geforce/dense_gemm.py b/examples/python/CuTeDSL/blackwell_geforce/dense_gemm.py index 74d2afa3..7fc7b421 100644 --- a/examples/python/CuTeDSL/blackwell_geforce/dense_gemm.py +++ b/examples/python/CuTeDSL/blackwell_geforce/dense_gemm.py @@ -29,7 +29,6 @@ import argparse from typing import Tuple, Type -import torch import cuda.bindings.driver as cuda import cutlass @@ -37,7 +36,6 @@ import cutlass.cute as cute import cutlass.cute.testing as testing import cutlass.utils as utils import cutlass.pipeline as pipeline -import cutlass.torch as cutlass_torch import cutlass.utils.hopper_helpers as sm90_utils """ @@ -49,15 +47,15 @@ using CUTE DSL. This GEMM kernel supports the following features: - Utilizes Tensor Memory Access (TMA) for efficient memory operations - - Utilizes non-Tensor Core MMA for matrix multiply-accumulate (MMA) operations + - Utilizes Blackwell MMA for matrix multiply-accumulate (MMA) operations - Supports multi-stage pipeline to overlap computation and memory access This GEMM works as follows: 1. Load A and B matrices from global memory (GMEM) to shared memory (SMEM) using TMA operations. -2. Perform matrix multiply-accumulate (MMA) operations using non-Tensor Core MMA instruction. +2. Perform matrix multiply-accumulate (MMA) operations using Blackwell MMA instruction. 3. Store results from registers (RMEM) to shared memory (SMEM), then to global memory (GMEM) with TMA operations. -Non-Tensor Core MMA instructions operate as follows: +Blackwell MMA instructions operate as follows: - Read matrix A from registers - Read matrix B from registers - Perform MMA operation and store the result in Accumulator(register) @@ -114,9 +112,7 @@ def parse_comma_separated_ints(s: str): def parse_arguments() -> argparse.Namespace: - parser = argparse.ArgumentParser( - description="Example of MxNxKxL GEMM on Blackwell Geforce." - ) + parser = argparse.ArgumentParser(description="Example of MxNxKxL GEMM on Blackwell Geforce.") parser.add_argument( "--mnkl", @@ -873,7 +869,10 @@ class Sm120GemmKernel: tRS_sD[(None, None, None, epi_buffer)], ) - cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) # barrier for sync self.epilog_sync_barrier.arrive_and_wait() @@ -1176,6 +1175,9 @@ def run( use_cold_l2: bool = False, **kwargs, ): + import torch + import cutlass.torch as cutlass_torch + print("Running Blackwell Geforce Dense GEMM with:") print(f"mnkl: {mnkl}") print( diff --git a/examples/python/CuTeDSL/cute/ffi/jit_argument.py b/examples/python/CuTeDSL/cute/ffi/jit_argument.py index c48cf8a1..58368551 100644 --- a/examples/python/CuTeDSL/cute/ffi/jit_argument.py +++ b/examples/python/CuTeDSL/cute/ffi/jit_argument.py @@ -240,10 +240,11 @@ import os import subprocess import shutil import tempfile -import torch def run_test(tmpdir=None, cmake_args="", cleanup=True): + import torch + try: current_dir = os.path.dirname(os.path.abspath(__file__)) diff --git a/examples/python/CuTeDSL/cute/torch_fake_tensor.py b/examples/python/CuTeDSL/cute/torch_fake_tensor.py index 37064441..60e936bd 100644 --- a/examples/python/CuTeDSL/cute/torch_fake_tensor.py +++ b/examples/python/CuTeDSL/cute/torch_fake_tensor.py @@ -27,8 +27,6 @@ # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. -import torch - import cutlass.cute as cute from cutlass.cute.runtime import from_dlpack @@ -66,6 +64,7 @@ def print_tensor(t: cute.Tensor): def run(): + import torch from torch._subclasses.fake_tensor import FakeTensorMode shape = (3, 4) diff --git a/examples/python/CuTeDSL/cute/tvm_ffi/aot_use_in_cpp_bundle.cpp b/examples/python/CuTeDSL/cute/tvm_ffi/aot_use_in_cpp_bundle.cpp index e1f0645b..37045630 100644 --- a/examples/python/CuTeDSL/cute/tvm_ffi/aot_use_in_cpp_bundle.cpp +++ b/examples/python/CuTeDSL/cute/tvm_ffi/aot_use_in_cpp_bundle.cpp @@ -1,6 +1,6 @@ // clang-format off /* - * SPDX-FileCopyrightText: Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: LicenseRef-NvidiaProprietary * * NVIDIA CORPORATION, its affiliates and licensors retain all intellectual diff --git a/examples/python/CuTeDSL/cute/tvm_ffi/aot_use_in_cpp_bundle.sh b/examples/python/CuTeDSL/cute/tvm_ffi/aot_use_in_cpp_bundle.sh index 3f110a4a..9c409724 100755 --- a/examples/python/CuTeDSL/cute/tvm_ffi/aot_use_in_cpp_bundle.sh +++ b/examples/python/CuTeDSL/cute/tvm_ffi/aot_use_in_cpp_bundle.sh @@ -28,7 +28,7 @@ #!/bin/bash # Set up library paths for runtime -export LD_LIBRARY_PATH=$(python3 -m cutlass.cute.export.aot_config --libdir):$(tvm-ffi-config --libdir) +export LD_LIBRARY_PATH=$(python3 -m cutlass.cute.export.aot_config --libdir):$(tvm-ffi-config --libdir):$LD_LIBRARY_PATH CUDA_HOME=/usr/local/cuda SOURCE_FILE="$(dirname "$0")/aot_use_in_cpp_bundle.cpp" diff --git a/examples/python/CuTeDSL/distributed/all_reduce_tma.py b/examples/python/CuTeDSL/distributed/all_reduce_tma.py index 5d6a4403..b3496c2b 100644 --- a/examples/python/CuTeDSL/distributed/all_reduce_tma.py +++ b/examples/python/CuTeDSL/distributed/all_reduce_tma.py @@ -53,8 +53,8 @@ Tile Assignment: - CTA i on rank r processes global_tile_id = r * ctas_per_rank + i TMA Usage Notes (for tutorial purposes, not perf-optimal): - - Uses TMALDG.1D to load from remote GPU memory via NVSHMEM addresses - - Uses TMASTG.1D to store to multicast address for broadcasting to all ranks + - Uses 1D TMA load to load from remote GPU memory via NVSHMEM addresses + - Uses 1D TMA load to store to multicast address for broadcasting to all ranks - Supports any input shape by flattening to 1D and tiling linearly - Pipeline with 2 stages overlaps TMA loads across ranks @@ -339,7 +339,7 @@ class AllReduceTmaKernel: producer_state.advance() # ====================================================================== - # Warp 1-4: Consumer - LDS, ADD, STS + # Warp 1-4: Consumer - Load from smem, ADD, Store to smem # ====================================================================== else: consumer_tid = tidx - self._tma_threads diff --git a/examples/python/CuTeDSL/experimental/blackwell/dense_block_scaled_gemm.py b/examples/python/CuTeDSL/experimental/blackwell/dense_block_scaled_gemm.py index 00bb784e..f5c2a542 100644 --- a/examples/python/CuTeDSL/experimental/blackwell/dense_block_scaled_gemm.py +++ b/examples/python/CuTeDSL/experimental/blackwell/dense_block_scaled_gemm.py @@ -213,12 +213,10 @@ class BlockScaledDenseGemmKernel: ) # UMMA ACC TMEM Layout - # ((MMA_M, MMA_N), REST_MMA_M, REST_MMA_N) - acc_shape = tiled_mma.partition_shape_C(mma_tiler_mnk[:2]) # ((MMA_M, MMA_N), REST_MMA_M, REST_MMA_N, ACC_STAGES) - tmem_accs_layout = tiled_mma.make_fragment_C( - cute.append(acc_shape, self.num_acc_stages) - ).layout + tmem_accs_layout = cute_ext.make_tmem_layout_acc( + tiled_mma, mma_tiler_mnk, self.num_acc_stages + ) sfa_tmem_layout = blockscaled_utils.make_tmem_layout_sfa( tiled_mma, @@ -318,21 +316,17 @@ class BlockScaledDenseGemmKernel: self.use_2cta_instrs, ) - # Performing layout calculations for one stage, in order to anticipate the - # required RMEM per thread and for reading from TMEM, and writing into SMEM - # tmem_acc: (MMA_M, MMA_N, MMA_REST_M, MMA_REST_N) - tmem_acc = tiled_mma.make_fragment_C(acc_shape) - # tmem_acc_epi: (EPI_TILE_M, EPI_TILE_N, EPI_REST_M, EPI_REST_N) - tmem_acc_epi = cute.flat_divide(tmem_acc[((None, None), 0, 0)], epi_tile) - tiled_copy_t2r = tcgen05.make_tmem_copy( - copy_atom_t2r, tmem_acc_epi[(None, None, 0, 0)] - ) + # Derive tiled_copy_t2r from the allocated TMEM buffer + accumulators = cute.zipped_divide(buffer_tmem_accs, ((epi_tile), 1)) + acc_epi_div = accumulators[((None, None), 0), 0] + tiled_copy_t2r = tcgen05.make_tmem_copy(copy_atom_t2r, acc_epi_div) thr_copy_t2r = tiled_copy_t2r.get_slice(tidx) - # gC_tile_epi: (EPI_TILE_M, EPI_TILE_N, EPI_REST_M, EPI_REST_N) + # Derive per-thread RMEM layout for the T2R epilogue copy gC_tile_epi = cute.flat_divide(gC_tile, epi_tile) - t2r_rmem_epi = thr_copy_t2r.partition_D(gC_tile_epi[(None, None, 0, 0)]) - acc_epi_rmem_layout = cute.make_fragment_like(t2r_rmem_epi.layout) + acc_epi_rmem_layout = cute_ext.make_t2r_rmem_layout( + tiled_copy_t2r, gC_tile_epi, tidx + ) # Allocate RMEM buffers buffer_rmem_t2r = cute_ext.allocate( diff --git a/examples/python/CuTeDSL/experimental/blackwell/dense_gemm.py b/examples/python/CuTeDSL/experimental/blackwell/dense_gemm.py index cd876904..b717cb76 100644 --- a/examples/python/CuTeDSL/experimental/blackwell/dense_gemm.py +++ b/examples/python/CuTeDSL/experimental/blackwell/dense_gemm.py @@ -463,18 +463,9 @@ class DenseGemmKernel: # - Has a capacity limit of 512 columns # - Requires specific layout patterns matching MMA instructions # - # partition_shape_C: Computes the accumulator shape based on MMA configuration. - # This returns the shape needed to store C = A × B results. - # - # cute.append(shape, stage): Appends a dimension for staging. - # For acc_stage=2: shape becomes (..., 2) for double-buffering. - # - # make_fragment_C: Creates a tensor descriptor with the appropriate layout - # for MMA accumulator storage. The .layout attribute extracts just the layout. - acc_shape = tiled_mma.partition_shape_C(mnk_tiler[:2]) # (M_tile, N_tile) - tmem_layout = tiled_mma.make_fragment_C( - cute.append(acc_shape, acc_stage) # Add stage dimension - ).layout + # make_tmem_layout_acc: Derives the TMEM accumulator buffer layout from the + # tiled MMA and MNK tiler, with the given number of pipeline stages. + tmem_layout = cute_ext.make_tmem_layout_acc(tiled_mma, mnk_tiler, acc_stage) # ======================================================================================== # STEP 10: ALLOCATE SMEM BUFFERS @@ -583,9 +574,6 @@ class DenseGemmKernel: # We derive the RMEM layout by partitioning the destination and extracting # the per-thread layout. # - # get_slice(tid_x): Gets the per-thread view of the tiled copy. - # partition_D: Partitions the destination tensor according to the copy layout. - # # CUTE ALGEBRA EXPLANATION - flat_divide: # --------------------------------------- # flat_divide(tensor, tiler) flattens all dimensions: @@ -594,20 +582,11 @@ class DenseGemmKernel: # Unlike zipped_divide which groups tile and rest separately, # flat_divide keeps everything flat, which is useful for iteration. # - # For epilogue: gC_mnl_epi = cute.flat_divide(gD_tile, epi_tile) - # This creates a flat view where we can iterate over sub-tiles with indices. - thr_copy_t2r = tiled_copy_t2r.get_slice(tid_x) + # make_t2r_rmem_layout: Derives the per-thread RMEM buffer layout + # produced by a TMEM->RMEM copy for a single epilogue iteration. gC_mnl_epi = cute.flat_divide(gD_tile, epi_tile) - - # Partition the output tensor according to the copy layout. - # tTR_gC has the thread's view of the output. - tTR_gC = thr_copy_t2r.partition_D(gC_mnl_epi) - - # make_fragment_like: Creates a layout matching a given tensor's layout. - # This is the standard way to derive RMEM layouts from copy partitions. - # The slicing [(None, None, None, 0, 0)] extracts one sub-tile's layout. - acc_d_rmem_layout = cute.make_fragment_like( - tTR_gC[(None, None, None, 0, 0)].layout + acc_d_rmem_layout = cute_ext.make_t2r_rmem_layout( + tiled_copy_t2r, gC_mnl_epi, tid_x ) # ======================================================================================== diff --git a/examples/python/CuTeDSL/experimental/blackwell/dense_gemm_2sm.py b/examples/python/CuTeDSL/experimental/blackwell/dense_gemm_2sm.py index 6273ad48..13561b0d 100644 --- a/examples/python/CuTeDSL/experimental/blackwell/dense_gemm_2sm.py +++ b/examples/python/CuTeDSL/experimental/blackwell/dense_gemm_2sm.py @@ -200,10 +200,9 @@ def sm100_4x4x1_kernel_builder( TMA_STORE_PIPE_DEPTH, ) - acc_shape = tiled_mma.partition_shape_C(mnk_tiler[:2]) - tmem_layout = tiled_mma.make_fragment_C( - cute.append(acc_shape, EPILOGUE_STAGE_DEPTH) - ).layout + tmem_layout = cute_ext.make_tmem_layout_acc( + tiled_mma, mnk_tiler, EPILOGUE_STAGE_DEPTH + ) bufferA = cute_ext.allocate( ab_dtype, @@ -251,11 +250,9 @@ def sm100_4x4x1_kernel_builder( tiled_copy_t2r = cute.nvgpu.tcgen05.make_tmem_copy(copy_atom_t2r, acc_epi_div) # Calculate the per thread destination size per iteration for output of TMEM and input of SMEM - thr_copy_t2r = tiled_copy_t2r.get_slice(tid_x) gC_mnl_epi = cute.flat_divide(tDgD, epi_tile) - tTR_gC = thr_copy_t2r.partition_D(gC_mnl_epi) - acc_d_rmem_layout = cute.make_fragment_like( - tTR_gC[(None, None, None, 0, 0)].layout + acc_d_rmem_layout = cute_ext.make_t2r_rmem_layout( + tiled_copy_t2r, gC_mnl_epi, tid_x ) bufferRAcc = cute_ext.allocate( diff --git a/examples/python/CuTeDSL/experimental/blackwell/dense_gemm_cute_pipeline.py b/examples/python/CuTeDSL/experimental/blackwell/dense_gemm_cute_pipeline.py index 588ee869..b4f7e551 100755 --- a/examples/python/CuTeDSL/experimental/blackwell/dense_gemm_cute_pipeline.py +++ b/examples/python/CuTeDSL/experimental/blackwell/dense_gemm_cute_pipeline.py @@ -1,4 +1,4 @@ -# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# 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 @@ -34,41 +34,28 @@ import cuda.bindings.driver as cuda import cutlass import cutlass.cute as cute import cutlass.cute.testing as testing -from cutlass.cute.runtime import from_dlpack import cutlass.utils as utils +from cutlass.utils import is_fp8_dtype, create_cute_tensor_for_fp8 import cutlass.pipeline as pipeline from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass import torch as cutlass_torch -import torch import cutlass.cute.experimental as cute_ext """ -A high-performance cluster launch control(CLC) dynamic persistent batched dense GEMM example -for the NVIDIA Blackwell SM100 architecture using CUTE DSL. - A high-performance persistent batched dense GEMM example for the NVIDIA Blackwell SM100 architecture using CUTE DSL. -This example attempts to show interoperability between cute.experimental and existing CUTE DSL APIs by using -cute.experimental APIs for TMA loading operations for A and B tensors. - -The CLC dynamic persistent scheduling technique performs dynamic loading balancing. -It has the ability to adapt available SMs rather than a statically selected number. To support this, -a new instruction is introduced to query for a new tile to compute. This new instruction is similar -to programmatic multicast in context of clusters in that the same starting tile ID for a given cluster -is broadcasted to all threadblocks in the cluster. -See `PTX documentation `. - - Matrix A is MxKxL, L is batch dimension, A can be row-major("K") or column-major("M") - Matrix B is NxKxL, L is batch dimension, B can be row-major("N") or column-major("K") - Matrix C is MxNxL, L is batch dimension, C can be row-major("N") or column-major("M") +This example attempts to show interoperability between cute.experimental and existing CUTE DSL APIs by using +cute.experimental APIs for TMA loading operations for A and B tensors. + This GEMM kernel supports the following features: - Utilizes Tensor Memory Access (TMA) for efficient memory operations - Utilizes Blackwell's tcgen05.mma for matrix multiply-accumulate (MMA) operations (including 2cta mma instructions) - Implements TMA multicast with cluster to reduce L2 memory traffic - Support persistent tile scheduling to better overlap memory load/store with mma between tiles - - Support CLC dynamic persistent tile scheduling to have near perfect load balancing - Support warp specialization to avoid explicit pipelining between mainloop load and mma This GEMM works as follows: @@ -207,6 +194,17 @@ class PersistentDenseGemmKernel: """This class implements batched matrix multiplication (C = A x B) with support for various data types and architectural features specific to Blackwell GPUs with persistent tile scheduling and warp specialization. + :param acc_dtype: Data type for accumulation during computation + :type acc_dtype: type[cutlass.Numeric] + :param use_2cta_instrs: Whether to use CTA group 2 for advanced thread cooperation + :type use_2cta_instrs: bool + :param mma_tiler_mn: Shape of the Matrix Multiply-Accumulate (MMA) tile (M,N) + :type mma_tiler_mn: Tuple[int, int] + :param cluster_shape_mn: Cluster dimensions (M,N) for parallel processing + :type cluster_shape_mn: Tuple[int, int] + :param use_tma_store: Whether to use Tensor Memory Access (TMA) for storing results + :type use_tma_store: bool + :note: In current version, A and B tensor must have the same data type - i.e., Float8E4M3FN for A and Float8E5M2 for B is not supported @@ -224,7 +222,7 @@ class PersistentDenseGemmKernel: :note: Supported C data types: - Float32 (for float32 and int32 accumulator data types) - Int32 (for float32 and int32 accumulator data types) - - Float16/BFloat16 (for fp32, fp16, and fp8 accumulator data types) + - Float16/BFloat16 (for fp16 and fp8 accumulator data types) - Int8/Uint8 (for uint8/int8 accumulator data types) - Float8E4M3FN/Float8E5M2 (for float32 accumulator data types) @@ -239,10 +237,9 @@ class PersistentDenseGemmKernel: acc_dtype=cutlass.Float32, use_2cta_instrs=True, mma_tiler_mn=(128, 128), - cluster_shape_mn=(2, 2), - use_tma_store=True + cluster_shape_mn=(2, 2) ) - gemm(a, b, c, max_active_clusters, stream, epilogue_op) + gemm(a, b, c, max_active_clusters, stream) """ def __init__( @@ -299,14 +296,8 @@ class PersistentDenseGemmKernel: self.epilogue_warp_id = (0, 1, 2, 3) self.mma_warp_id = 4 self.tma_warp_id = 5 - self.sched_warp_id = 6 self.threads_per_cta = 32 * len( - ( - self.mma_warp_id, - self.tma_warp_id, - self.sched_warp_id, - *self.epilogue_warp_id, - ) + (self.mma_warp_id, self.tma_warp_id, *self.epilogue_warp_id) ) # Set barrier id for cta sync, epilogue sync and tmem ptr sync self.epilog_sync_bar_id = 1 @@ -341,12 +332,12 @@ class PersistentDenseGemmKernel: tiled_mma = self._create_tiled_mma() # Compute mma/cluster/tile shapes - self.mma_inst_shape_k = cute.size(tiled_mma.shape_mnk, mode=[2]) + mma_inst_shape_k = cute.size(tiled_mma.shape_mnk, mode=[2]) mma_inst_tile_k = 4 self.mma_tiler = ( self.mma_tiler[0], self.mma_tiler[1], - self.mma_inst_shape_k * mma_inst_tile_k, + mma_inst_shape_k * mma_inst_tile_k, ) self.cta_tile_shape_mnk = ( self.mma_tiler[0] // cute.size(tiled_mma.thr_id.shape), @@ -398,9 +389,6 @@ class PersistentDenseGemmKernel: c_smem_layout, ) - # Setup clc stage by default - self.num_clc_stage = 1 - # Compute A/B/C shared memory layout self.a_smem_layout_staged = utils.sm100.make_smem_layout_a( tiled_mma, self.mma_tiler, self.a_dtype, self.num_ab_stage @@ -479,9 +467,7 @@ class PersistentDenseGemmKernel: a_copy_size = cute.size_in_bytes(self.a_dtype, a_smem_layout) b_copy_size = cute.size_in_bytes(self.b_dtype, b_smem_layout) - self.num_tma_load_bytes = (b_copy_size + a_copy_size) * atom_thr_size - # Response size is 4B * 4 elements - self.num_clc_response_bytes = 16 + self.num_tma_load_bytes = (a_copy_size + b_copy_size) * atom_thr_size # Setup TMA store for C tma_atom_c = None @@ -494,9 +480,7 @@ class PersistentDenseGemmKernel: # Compute grid size self.tile_sched_params, grid = self._compute_grid( - c, - self.cta_tile_shape_mnk, - self.cluster_shape_mn, + c, self.cta_tile_shape_mnk, self.cluster_shape_mn, max_active_clusters ) # Launch the kernel synchronously @@ -535,7 +519,7 @@ class PersistentDenseGemmKernel: b_smem_layout_staged: cute.ComposedLayout, c_smem_layout_staged: Union[cute.Layout, cute.ComposedLayout, None], epi_tile: cute.Tile, - tile_sched_params: utils.ClcDynamicPersistentTileSchedulerParams, + tile_sched_params: utils.PersistentTileSchedulerParams, epilogue_op: cutlass.Constexpr, ): """ @@ -563,19 +547,29 @@ class PersistentDenseGemmKernel: cta_rank_in_cluster = cute.arch.make_warp_uniform( cute.arch.block_idx_in_cluster() ) - is_first_cta_in_cluster = cta_rank_in_cluster == 0 block_in_cluster_coord_vmnk = cluster_layout_vmnk.get_flat_coord( cta_rank_in_cluster ) + cluster_layout_v_size = cute.size(cluster_layout_vmnk.shape[0]) # Coord inside cta tidx, _, _ = cute.arch.thread_idx() - ## Tiling the global tensors for LIR TMA Loads - tiler_mk = (self.mma_tiler[0], self.mma_tiler[2]) - gA = cute.zipped_divide(mA, tiler_mk) + ## Tiling the global tensors for cute.experimental TMA Loads + num_mma_ctas = cute.size(tiled_mma.thr_id.shape) + cta_tile_shape_mnk = cute.shape_div(self.mma_tiler, (num_mma_ctas, 1, 1)) + # A is tiled (M/2, K) for 2CTA + a_tiler_mk = (cta_tile_shape_mnk[0], cta_tile_shape_mnk[2]) + # B is tiled (N/2, K) for 2CTA + b_tiler_nk = (cta_tile_shape_mnk[1] // num_mma_ctas, cta_tile_shape_mnk[2]) - tiler_nk = (self.mma_tiler[1], self.mma_tiler[2]) - gB = cute.zipped_divide(mB, tiler_nk) + gA = cute.zipped_divide(mA, a_tiler_mk) + gB = cute.zipped_divide(mB, b_tiler_nk) + + # Determine pipeline operation types based on 2-CTA mode and TMA multicast + if cutlass.const_expr(self.use_2cta_instrs): + tma_operation_type = cute_ext.OperationTypeEnum.SM100_TMA_LOAD_2SM + else: + tma_operation_type = cute_ext.OperationTypeEnum.SM90_TMA_LOAD # # Alloc and init: a+b full/empty, accumulator full/empty, tensor memory dealloc barrier @@ -589,8 +583,6 @@ class PersistentDenseGemmKernel: ] tmem_dealloc_mbar_ptr: cutlass.Int64 tmem_holding_buf: cutlass.Int32 - clc_ptr: cute.struct.MemRange[cutlass.Int64, self.num_clc_stage * 2] - clc_response_ptr: cute.struct.MemRange[cutlass.Int32, 1] smem = utils.SmemAllocator() storage = smem.allocate(SharedStorage) @@ -628,42 +620,6 @@ class PersistentDenseGemmKernel: defer_sync=True, ) - # Initialize clc_pipeline (barrier) and states - clc_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) - cluster_size = cute.size(self.cluster_shape_mn) - num_clc_consumer_threads = 32 * len( - ( - self.sched_warp_id, - *( - cluster_size - * ( - self.mma_warp_id, - self.tma_warp_id, - *self.epilogue_warp_id, - ) - ), - ) - ) - - clc_pipeline_consumer_group = pipeline.CooperativeGroup( - pipeline.Agent.Thread, num_clc_consumer_threads - ) - clc_pipeline = pipeline.PipelineClcFetchAsync.create( - barrier_storage=storage.clc_ptr.data_ptr(), - num_stages=self.num_clc_stage, - producer_group=clc_pipeline_producer_group, - consumer_group=clc_pipeline_consumer_group, - tx_count=self.num_clc_response_bytes, - cta_layout_vmnk=cluster_layout_vmnk, - defer_sync=True, - ) - # Initial clc response pointer - clc_response_ptr = storage.clc_response_ptr.data_ptr() - - clc_consumer_state = pipeline.make_pipeline_state( - pipeline.PipelineUserType.Consumer, self.num_clc_stage - ) - tmem_alloc_barrier = pipeline.NamedBarrier( barrier_id=self.tmem_alloc_sync_bar_id, num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id)), @@ -684,7 +640,7 @@ class PersistentDenseGemmKernel: ) # Cluster arrive after barrier init - pipeline_init_arrive(cluster_shape_mn=self.cluster_shape_mn, is_relaxed=True) + pipeline_init_arrive(cluster_shape_mn=cluster_layout_vmnk, is_relaxed=True) # # Setup smem tensor A/B/C @@ -740,10 +696,8 @@ class PersistentDenseGemmKernel: # # Partition shared/tensor memory tensor for TiledMMA_A/B/C # - # (MMA, MMA_M, MMA_K, STAGE) tCrA = tiled_mma.make_fragment_A(bufferA) - # (MMA, MMA_N, MMA_K, STAGE) tCrB = tiled_mma.make_fragment_B(bufferB) # (MMA, MMA_M, MMA_N) @@ -756,16 +710,15 @@ class PersistentDenseGemmKernel: # # Cluster wait before tensor memory alloc # - pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mn) + pipeline_init_wait(cluster_shape_mn=cluster_layout_vmnk) # # Construct the scheduler # - tile_sched = utils.ClcDynamicPersistentTileScheduler.create( + tile_sched = utils.StaticPersistentTileScheduler.create( tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim(), - clc_response_ptr, ) work_tile = tile_sched.initial_work_tile_info() @@ -777,25 +730,28 @@ class PersistentDenseGemmKernel: # # Persistent tile scheduling loop # + while work_tile.is_valid_tile: # Get tile coord from tile scheduler cur_tile_coord = work_tile.tile_idx - mma_tile_coord_mnl = ( - cur_tile_coord[0] // cute.size(tiled_mma.thr_id.shape), - cur_tile_coord[1], - cur_tile_coord[2], - ) # # Slice to per mma tile index # - gA_tile = gA[ - (None, None), (mma_tile_coord_mnl[0], None, mma_tile_coord_mnl[2]) - ] + gA_tile = gA[(None, None), (cur_tile_coord[0], None, cur_tile_coord[2])] - gB_tile = gB[ - (None, None), (mma_tile_coord_mnl[1], None, mma_tile_coord_mnl[2]) - ] + # For B loading in 2-CTA mode, compute proper N coordinate + if cutlass.const_expr(self.use_2cta_instrs): + # In 2CTA mode, the cur_tile_coord[1] gives a full MMA tile, but we want a CTA level tile to load + # Each CTA in the pair loads a half of the N tile + gB_tma_coord_n = ( + cluster_layout_v_size * cur_tile_coord[1] + + bidx % cluster_layout_v_size + ) + else: + gB_tma_coord_n = cur_tile_coord[1] + + gB_tile = gB[(None, None), (gB_tma_coord_n, None, cur_tile_coord[2])] # Peek (try_wait) AB buffer empty for k_tile = prefetch_k_tile_cnt ab_producer.reset() @@ -811,7 +767,6 @@ class PersistentDenseGemmKernel: handle = ab_producer.acquire_and_advance(peek_ab_empty_status) idx = handle.index bufferA_sliced = bufferA[None, None, None, idx] - # print("*********Type of bufferA_sliced: ", bufferA_sliced.type) a_cta_v_map = cute_ext.get_cta_v_map_ab( mA, self.mma_tiler, tiled_mma, "A" ) @@ -827,6 +782,7 @@ class PersistentDenseGemmKernel: handle.barrier.value, cta_v_map=a_cta_v_map, update_expect_tx=False, # Does not automatically update the mbarrier's transaction bytes + tma_operation_type=tma_operation_type, ) cute_ext.tma_load( @@ -835,6 +791,7 @@ class PersistentDenseGemmKernel: handle.barrier.value, cta_v_map=b_cta_v_map, update_expect_tx=False, # Does not automatically update the mbarrier's transaction bytes + tma_operation_type=tma_operation_type, ) # Peek (try_wait) AB buffer empty for k_tile = prefetch_k_tile_cnt + k_tile + 1 @@ -845,42 +802,14 @@ class PersistentDenseGemmKernel: # # Advance to next tile # - clc_pipeline.consumer_wait(clc_consumer_state) + tile_sched.advance_to_next_work() work_tile = tile_sched.get_current_work() - clc_pipeline.consumer_release(clc_consumer_state) - clc_consumer_state.advance() + # # Wait A/B buffer empty # ab_producer.tail() - # - # Sched warp - # - - if warp_idx == self.sched_warp_id and is_first_cta_in_cluster: - # - # Persistent tile scheduling loop - # - clc_producer_state = pipeline.make_pipeline_state( - pipeline.PipelineUserType.ProducerConsumer, self.num_clc_stage - ) - - while work_tile.is_valid_tile: - # - # Advance to next tile - # - clc_pipeline.producer_acquire(clc_producer_state) - mbarrier_addr = clc_pipeline.producer_get_barrier(clc_producer_state) - tile_sched.advance_to_next_work(mbarrier_addr) - clc_producer_state.advance() - - clc_pipeline.consumer_wait(clc_consumer_state) - work_tile = tile_sched.get_current_work() - clc_pipeline.consumer_release(clc_consumer_state) - clc_consumer_state.advance() - clc_pipeline.producer_tail(clc_producer_state) - # # Specialized MMA warp # @@ -896,6 +825,7 @@ class PersistentDenseGemmKernel: # # Persistent tile scheduling loop # + acc_producer_state = pipeline.make_pipeline_state( pipeline.PipelineUserType.Producer, self.num_acc_stage ) @@ -925,6 +855,11 @@ class PersistentDenseGemmKernel: if is_leader_cta: acc_pipeline.producer_acquire(acc_producer_state) + # + # Reset the ACCUMULATE field for each tile + # + tiled_mma.set(tcgen05.Field.ACCUMULATE, False) + # # Mma mainloop # @@ -932,12 +867,21 @@ class PersistentDenseGemmKernel: if is_leader_cta: # Conditionally wait for AB buffer full handle = ab_consumer.wait_and_advance(peek_ab_full_status) + # tCtAcc += tCrA * tCrB - tiled_mma.set(tcgen05.Field.ACCUMULATE, k_tile != 0) - tile_crd = (None, None, None, handle.index) - cute.gemm( - tiled_mma, tCtAcc, tCrA[tile_crd], tCrB[tile_crd], tCtAcc - ) + num_kblocks = cute.size(tCrA, mode=[2]) + for kblk_idx in cutlass.range(num_kblocks, unroll_full=True): + kblk_crd = (None, None, kblk_idx, handle.index) + + cute.gemm( + tiled_mma, + tCtAcc, + tCrA[kblk_crd], + tCrB[kblk_crd], + tCtAcc, + ) + # Enable accumulate on tCtAcc after first kblock + tiled_mma.set(tcgen05.Field.ACCUMULATE, True) # Async arrive AB buffer empty handle.release() @@ -957,10 +901,9 @@ class PersistentDenseGemmKernel: # # Advance to next tile # - clc_pipeline.consumer_wait(clc_consumer_state) + tile_sched.advance_to_next_work() work_tile = tile_sched.get_current_work() - clc_pipeline.consumer_release(clc_consumer_state) - clc_consumer_state.advance() + # # Wait for accumulator buffer empty # @@ -1017,6 +960,12 @@ class PersistentDenseGemmKernel: cur_tile_coord[1], cur_tile_coord[2], ) + # + # Pre-advance to next tile + # + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + num_tiles_executed = tile_sched.num_tiles_executed if cutlass.const_expr(self.use_tma_store): acc_consumer_state = utils.gemm.sm100.epilogue_tma_store( @@ -1047,13 +996,6 @@ class PersistentDenseGemmKernel: acc_consumer_state, acc_pipeline, ) - # - # Advance to next tile - # - clc_pipeline.consumer_wait(clc_consumer_state) - work_tile = tile_sched.get_current_work() - clc_pipeline.consumer_release(clc_consumer_state) - clc_consumer_state.advance() if cutlass.const_expr(self.use_tma_store): # Wait for C store complete @@ -1068,13 +1010,13 @@ class PersistentDenseGemmKernel: tmem.relinquish_alloc_permit() tmem.free(tmem_ptr) - @staticmethod def _compute_grid( c: cute.Tensor, cta_tile_shape_mnk: Tuple[int, int, int], cluster_shape_mn: Tuple[int, int], - ) -> Tuple[utils.ClcDynamicPersistentTileSchedulerParams, Tuple[int, int, int]]: + max_active_clusters: cutlass.Constexpr, + ) -> Tuple[utils.PersistentTileSchedulerParams, Tuple[int, int, int]]: """Use persistent tile scheduler to compute the grid size for the output tensor C. :param c: The output tensor C @@ -1083,21 +1025,25 @@ class PersistentDenseGemmKernel: :type cta_tile_shape_mnk: tuple[int, int, int] :param cluster_shape_mn: Shape of each cluster in M, N dimensions. :type cluster_shape_mn: tuple[int, int] + :param max_active_clusters: Maximum number of active clusters. + :type max_active_clusters: cutlass.Constexpr :return: A tuple containing: - tile_sched_params: Parameters for the persistent tile scheduler. - grid: Grid shape for kernel launch. - :rtype: Tuple[utils.ClcDynamicPersistentTileSchedulerParams, tuple[int, int, int]] + :rtype: Tuple[utils.PersistentTileSchedulerParams, tuple[int, int, int]] """ c_shape = cute.slice_(cta_tile_shape_mnk, (None, None, 0)) gc = cute.zipped_divide(c, tiler=c_shape) num_ctas_mnl = gc[(0, (None, None, None))].shape cluster_shape_mnl = (*cluster_shape_mn, 1) - tile_sched_params = utils.ClcDynamicPersistentTileSchedulerParams( + tile_sched_params = utils.PersistentTileSchedulerParams( num_ctas_mnl, cluster_shape_mnl ) - grid = utils.ClcDynamicPersistentTileScheduler.get_grid_shape(tile_sched_params) + grid = utils.StaticPersistentTileScheduler.get_grid_shape( + tile_sched_params, max_active_clusters + ) return tile_sched_params, grid @@ -1128,13 +1074,18 @@ class PersistentDenseGemmKernel: return num_tmem_alloc_cols def check_supported_dtypes( - self, ab_dtype: Type[cutlass.Numeric], c_dtype: Type[cutlass.Numeric] - ) -> bool: + self, + a_dtype: Type[cutlass.Numeric], + b_dtype: Type[cutlass.Numeric], + c_dtype: Type[cutlass.Numeric], + ): """ Check if the dtypes are valid - :param ab_dtype: The data type of the A and B operands - :type ab_dtype: Type[cutlass.Numeric] + :param a_dtype: The data type of the A operands + :type a_dtype: Type[cutlass.Numeric] + :param b_dtype: The data type of the B operands + :type b_dtype: Type[cutlass.Numeric] :param acc_dtype: The data type of the accumulator :type acc_dtype: Type[cutlass.Numeric] :param c_dtype: The data type of the output tensor @@ -1151,8 +1102,10 @@ class PersistentDenseGemmKernel: cutlass.Float8E4M3FN, cutlass.Float8E5M2, } - if ab_dtype not in valid_ab_dtypes: - raise testing.CantImplementError(f"Unsupported AB dtype: {ab_dtype}") + if a_dtype not in valid_ab_dtypes or b_dtype not in valid_ab_dtypes: + raise testing.CantImplementError( + f"Unsupported AB dtype: {a_dtype} and {b_dtype}" + ) if self.acc_dtype not in {cutlass.Float32, cutlass.Float16, cutlass.Int32}: raise testing.CantImplementError( @@ -1176,8 +1129,13 @@ class PersistentDenseGemmKernel: cutlass.Int32: {cutlass.Uint8, cutlass.Int8}, } # Check compatibility between accumulator type and AB type - if ab_dtype not in acc_ab_compatibility[self.acc_dtype]: - return False + if ( + a_dtype not in acc_ab_compatibility[self.acc_dtype] + or b_dtype not in acc_ab_compatibility[self.acc_dtype] + ): + raise testing.CantImplementError( + f"Unsupported AB dtype: {a_dtype} and {b_dtype} for accumulator dtype: {self.acc_dtype}" + ) # Define compatibility mapping between accumulator type and C type acc_c_compatibility = { @@ -1206,20 +1164,15 @@ class PersistentDenseGemmKernel: } # Check compatibility between accumulator type and C type if c_dtype not in acc_c_compatibility[self.acc_dtype]: - return False + raise testing.CantImplementError( + f"Unsupported C dtype: {c_dtype} for accumulator dtype: {self.acc_dtype}" + ) - return True - - def check_mma_tiler_and_cluster_shape(self) -> bool: + def check_mma_tiler_and_cluster_shape(self): """Check if the mma tiler and cluster shape are valid. :raises testing.CantImplementError: If the mma tiler and cluster shape are invalid """ - # This kernel does not support 2CTA MMA instructions - if self.use_2cta_instrs: - raise testing.CantImplementError( - f"2 CTA instructions are not supported by this kernel {self.use_2cta_instrs}" - ) # Skip invalid mma tile shape if not ( (not self.use_2cta_instrs and self.mma_tiler_mn[0] in [64, 128]) @@ -1256,12 +1209,13 @@ class PersistentDenseGemmKernel: n: int, k: int, l: int, - ab_dtype: Type[cutlass.Numeric], + a_dtype: Type[cutlass.Numeric], + b_dtype: Type[cutlass.Numeric], c_dtype: Type[cutlass.Numeric], a_major: str, b_major: str, c_major: str, - ) -> bool: + ): """ Check if the tensor alignment is valid @@ -1273,8 +1227,10 @@ class PersistentDenseGemmKernel: :type k: int :param l: The number of columns in the C tensor :type l: int - :param ab_dtype: The data type of the A and B operands - :type ab_dtype: Type[cutlass.Numeric] + :param a_dtype: The data type of the A operands + :type a_dtype: Type[cutlass.Numeric] + :param b_dtype: The data type of the B operands + :type b_dtype: Type[cutlass.Numeric] :param c_dtype: The data type of the output tensor :type c_dtype: Type[cutlass.Numeric] :param a_major: The major axis of the A tensor @@ -1284,8 +1240,7 @@ class PersistentDenseGemmKernel: :param c_major: The major axis of the C tensor :type c_major: str - :return: True if the problem shape is valid, False otherwise - :rtype: bool + :raises testing.CantImplementError: If the tensor alignment is invalid """ # TODO: move to utils @@ -1296,15 +1251,15 @@ class PersistentDenseGemmKernel: return num_major_elements % num_contiguous_elements == 0 if ( - not check_contiguous_16B_alignment(ab_dtype, a_major == "m", (m, k, l)) - or not check_contiguous_16B_alignment(ab_dtype, b_major == "n", (n, k, l)) + not check_contiguous_16B_alignment(a_dtype, a_major == "m", (m, k, l)) + or not check_contiguous_16B_alignment(b_dtype, b_major == "n", (n, k, l)) or not check_contiguous_16B_alignment(c_dtype, c_major == "m", (m, n, l)) ): raise testing.CantImplementError( - f"Invalid tensor alignment: {m}, {n}, {k}, {l}, {ab_dtype}, {c_dtype}, {a_major}, {b_major}, {c_major}" + f"Invalid tensor alignment: {m}, {n}, {k}, {l}, {a_dtype}, {b_dtype}, {c_dtype}, {a_major}, {b_major}, {c_major}" ) - def check_epilog_store_option(self, m: int, n: int) -> bool: + def check_epilog_store_option(self, m: int, n: int): """ Check if the epilogue store option is valid @@ -1329,7 +1284,8 @@ class PersistentDenseGemmKernel: def can_implement( self, mnkl: Tuple[int, int, int, int], - ab_dtype: Type[cutlass.Numeric], + a_dtype: Type[cutlass.Numeric], + b_dtype: Type[cutlass.Numeric], c_dtype: Type[cutlass.Numeric], a_major: str, b_major: str, @@ -1340,8 +1296,10 @@ class PersistentDenseGemmKernel: :param mnkl: Problem size as a tuple (M, N, K, L). :type mnkl: Tuple[int, int, int, int] - :param ab_dtype: Data type for input tensors A and B. - :type ab_dtype: Type[cutlass.Numeric] + :param a_dtype: Data type for input tensors A. + :type a_dtype: Type[cutlass.Numeric] + :param b_dtype: Data type for input tensors B. + :type b_dtype: Type[cutlass.Numeric] :param c_dtype: Data type for output tensor C. :type c_dtype: Type[cutlass.Numeric] :param a_major: Major dimension of the A tensor layout ("m" or "k"). @@ -1356,14 +1314,14 @@ class PersistentDenseGemmKernel: try: # Skip unsupported types - self.check_supported_dtypes(ab_dtype, c_dtype) + self.check_supported_dtypes(a_dtype, b_dtype, c_dtype) # Skip invalid mma tile shape and cluster shape self.check_mma_tiler_and_cluster_shape() m, n, k, l = mnkl self.check_tensor_alignment( - m, n, k, l, ab_dtype, c_dtype, a_major, b_major, c_major + m, n, k, l, a_dtype, b_dtype, c_dtype, a_major, b_major, c_major ) self.check_epilog_store_option(m, n) except testing.CantImplementError: @@ -1371,6 +1329,119 @@ class PersistentDenseGemmKernel: return True +@cute.experimental.jit +def bmm( + gemm_op: cutlass.Constexpr, + a: cute.Tensor, # (l, m, k) + b: cute.Tensor, # (l, k, n) + c: cute.Tensor, # (l, m, n) + max_active_clusters: cutlass.Constexpr, + stream: cuda.CUstream, + epilogue_op: cutlass.Constexpr = lambda x: x, +): + """ + Wrapper API for persistent GEMM kernel to follow the convention of PyTorch's batch matrix-multiply (bmm). + + Internally, the tensors are permuted to match CuTe's convention: + - a: (m, k, l) + - b: (n, k, l) + - c: (m, n, l) + + :param gemm_op: Kernel operation, expects (a, b, c, max_active_clusters, stream, epilogue_op) + :type gemm_op: cutlass.Constexpr + :param a: Input tensor of shape (l, m, k) + :type a: cute.Tensor + :param b: Input tensor of shape (l, k, n) + :type b: cute.Tensor + :param c: Output tensor of shape (l, m, n) + :type c: cute.Tensor + :param max_active_clusters: Maximum number of hardware clusters to launch + :type max_active_clusters: cutlass.Constexpr + :param epilogue_op: Optional elementwise lambda function to apply per output element, defaults to identity + :type epilogue_op: cutlass.Constexpr, optional + """ + # (l,m,k) -> (m,k,l) + a = cute.make_tensor(a.iterator, cute.select(a.layout, mode=[1, 2, 0])) + # (l,k,n) -> (n,k,l) + b = cute.make_tensor(b.iterator, cute.select(b.layout, mode=[2, 1, 0])) + # (l,m,n) -> (m,n,l) + c = cute.make_tensor(c.iterator, cute.select(c.layout, mode=[1, 2, 0])) + + gemm_op(a, b, c, max_active_clusters, stream, epilogue_op) + + +@lru_cache(maxsize=1) +def prepare_tensors( + mnkl: Tuple[int, int, int, int], + a_dtype: Type[cutlass.Numeric], + b_dtype: Type[cutlass.Numeric], + c_dtype: Type[cutlass.Numeric], + a_major: str, + b_major: str, + c_major: str, + init_random: bool = True, + normal_mean: float = 0.0, + normal_std: float = 1.0, +): + """Prepare tensors for GEMM. + + Returns: + Tuple of (a_f32, b_f32, c_f32, a_storage, b_storage, c_storage): + - *_f32: Float32 tensors with the logical data (for reference and fp8 conversion) + - *_storage: Storage tensors for DLPack (uint8 for fp8, otherwise the target dtype) + """ + import torch + from cutlass.torch import dtype as torch_dtype + + m, n, k, l = mnkl + + if a_major == "k": + a_f32 = torch.empty((l, m, k), dtype=torch.float32, device="cuda") + elif a_major == "m": + a_f32 = torch.empty((l, k, m), dtype=torch.float32, device="cuda").permute( + 0, 2, 1 + ) + + if b_major == "n": + b_f32 = torch.empty((l, k, n), dtype=torch.float32, device="cuda") + elif b_major == "k": + b_f32 = torch.empty((l, n, k), dtype=torch.float32, device="cuda").permute( + 0, 2, 1 + ) + + if c_major == "n": + c_f32 = torch.empty((l, m, n), dtype=torch.float32, device="cuda") + elif c_major == "m": + c_f32 = torch.empty((l, n, m), dtype=torch.float32, device="cuda").permute( + 0, 2, 1 + ) + + if init_random: + # Uniform random initialization in range [-2, 3) + a_f32.random_(-2, 3) + b_f32.random_(-2, 3) + c_f32.random_(-2, 3) + + else: + # Normal (Gaussian) initialization with user-specified mean and std + a_f32.normal_(mean=normal_mean, std=normal_std) + b_f32.normal_(mean=normal_mean, std=normal_std) + c_f32.normal_(mean=normal_mean, std=normal_std) + + # For float8 types, use uint8 as storage type to avoid dlpack limitation + # (dlpack doesn't support float8 types) + # For other types, convert to the target dtype + a_storage_dtype = torch.uint8 if is_fp8_dtype(a_dtype) else torch_dtype(a_dtype) + b_storage_dtype = torch.uint8 if is_fp8_dtype(b_dtype) else torch_dtype(b_dtype) + c_storage_dtype = torch.uint8 if is_fp8_dtype(c_dtype) else torch_dtype(c_dtype) + + a_storage = a_f32.to(dtype=a_storage_dtype) + b_storage = b_f32.to(dtype=b_storage_dtype) + c_storage = c_f32.to(dtype=c_storage_dtype) + + return (a_f32, b_f32, c_f32, a_storage, b_storage, c_storage) + + @lru_cache(maxsize=1) def compile_bmm( mnkl: Tuple[int, int, int, int], @@ -1381,10 +1452,10 @@ def compile_bmm( a_major: str, b_major: str, c_major: str, - mma_tiler: Union[Tuple[int, int], Tuple[int, int, int]] = (128, 128), - cluster_shape_mn: Tuple[int, int] = (1, 1), + mma_tiler_mn: Tuple[int, int] = (256, 256), + cluster_shape_mn: Tuple[int, int] = (2, 1), max_active_clusters: cutlass.Constexpr = None, - use_2cta_instrs: bool = False, + use_2cta_instrs: bool = True, use_tma_store: bool = True, epilogue_op: cutlass.Constexpr = lambda x: x, ): @@ -1393,36 +1464,23 @@ def compile_bmm( gemm = PersistentDenseGemmKernel( acc_dtype, use_2cta_instrs, - mma_tiler, + mma_tiler_mn, cluster_shape_mn, use_tma_store, ) - # Check if configuration can be implemented can_implement = gemm.can_implement( - mnkl, a.element_type, c.element_type, a_major, b_major, c_major + mnkl, a.element_type, b.element_type, c.element_type, a_major, b_major, c_major ) if not can_implement: raise testing.CantImplementError( f"The current config which is invalid/unsupported: use_2cta_instrs = {use_2cta_instrs}, " - f"mma_tiler = {mma_tiler}, cluster_shape_mn = {cluster_shape_mn}, " + f"mma_tiler_mn = {mma_tiler_mn}, cluster_shape_mn = {cluster_shape_mn}, " f"use_tma_store = {use_tma_store}" ) stream = make_fake_stream() - return cute.compile(gemm, a, b, c, max_active_clusters, stream, epilogue_op) - - -def compare_reference(a_torch_cpu, b_torch_cpu, c_torch_gpu, c_dtype, tolerance): - ref = torch.einsum("mkl,nkl->mnl", a_torch_cpu, b_torch_cpu) - - _, ref_torch_gpu = cutlass_torch.cute_tensor_like( - ref, c_dtype, is_dynamic_layout=True, assumed_align=16 - ) - ref_result = ref_torch_gpu.cpu() - torch.testing.assert_close( - c_torch_gpu.cpu(), ref_result, atol=tolerance, rtol=1e-05 - ) + return cute.compile(bmm, gemm, a, b, c, max_active_clusters, stream, epilogue_op) def run( @@ -1433,9 +1491,9 @@ def run( a_major: str, b_major: str, c_major: str, - mma_tiler_mn: Tuple[int, int] = (128, 128), - cluster_shape_mn: Tuple[int, int] = (1, 1), - use_2cta_instrs: bool = False, + mma_tiler_mn: Tuple[int, int] = (256, 256), + cluster_shape_mn: Tuple[int, int] = (2, 1), + use_2cta_instrs: bool = True, use_tma_store: bool = True, tolerance: float = 1e-01, warmup_iterations: int = 0, @@ -1490,6 +1548,7 @@ def run( :return: Execution time of the GEMM kernel. :rtype: float """ + import torch from cutlass.torch import dtype as torch_dtype if not torch.cuda.is_available(): @@ -1506,9 +1565,23 @@ def run( ) # Run and verify BMM with torch - m, n, k, l = mnkl - a_tensor, b_tensor, c_tensor, a_torch_cpu, b_torch_cpu, c_torch_cpu, c_torch_gpu = ( - create_tensors(l, m, n, k, a_major, b_major, c_major, ab_dtype, c_dtype) + a_f32, b_f32, c_f32, a_storage, b_storage, c_storage = prepare_tensors( + mnkl, ab_dtype, ab_dtype, c_dtype, a_major, b_major, c_major + ) + + leading_dim_a = 2 if a_major == "k" else 1 + leading_dim_b = 1 if b_major == "k" else 2 + leading_dim_c = 2 if c_major == "n" else 1 + + # Create CuTe tensors, passing float32 source for fp8 conversion + a_tensor = create_cute_tensor_for_fp8( + a_storage, ab_dtype, leading_dim_a, source_f32_tensor=a_f32 + ) + b_tensor = create_cute_tensor_for_fp8( + b_storage, ab_dtype, leading_dim_b, source_f32_tensor=b_f32 + ) + c_tensor = create_cute_tensor_for_fp8( + c_storage, c_dtype, leading_dim_c, source_f32_tensor=c_f32 ) compiled_fn = compile_bmm( @@ -1541,37 +1614,54 @@ def run( compiled_fn(a_tensor, b_tensor, c_tensor, current_stream) # Manually quantize to be comparable - compare_reference(a_torch_cpu, b_torch_cpu, c_torch_gpu, c_dtype, tolerance) + # Use float32 source data for reference calculation + ref = ( + torch.bmm(a_f32, b_f32) + .to(dtype=torch_dtype(c_dtype)) + .to(dtype=torch.float32) + ) + # Read back the result from CuTe tensor (c_storage was updated in-place) + torch.testing.assert_close( + c_storage.to(dtype=torch.float32), ref, atol=tolerance, rtol=1e-03 + ) if not benchmark: return 0 def generate_tensors(): - init_normal = ab_dtype not in [cutlass.Int8, cutlass.Uint8] - ( - a_tensor, - b_tensor, - c_tensor, - a_torch_cpu, - b_torch_cpu, - c_torch_cpu, - c_torch_gpu, - ) = create_tensors(l, m, n, k, a_major, b_major, c_major, ab_dtype, c_dtype) + a_f32, b_f32, c_f32, a_st, b_st, c_st = prepare_tensors( + mnkl, + ab_dtype, + ab_dtype, + c_dtype, + a_major, + b_major, + c_major, + ) + a_tensor = create_cute_tensor_for_fp8( + a_st, ab_dtype, leading_dim_a, source_f32_tensor=a_f32 + ) + b_tensor = create_cute_tensor_for_fp8( + b_st, ab_dtype, leading_dim_b, source_f32_tensor=b_f32 + ) + c_tensor = create_cute_tensor_for_fp8( + c_st, c_dtype, leading_dim_c, source_f32_tensor=c_f32 + ) return testing.JitArguments(a_tensor, b_tensor, c_tensor, current_stream) workspace_count = 1 if use_cold_l2: one_workspace_bytes = ( - a_torch_cpu.numel() * a_torch_cpu.element_size() - + b_torch_cpu.numel() * b_torch_cpu.element_size() - + c_torch_cpu.numel() * c_torch_cpu.element_size() + a_storage.numel() * a_storage.element_size() + + b_storage.numel() * b_storage.element_size() + + c_storage.numel() * c_storage.element_size() ) workspace_count = testing.get_workspace_count( one_workspace_bytes, warmup_iterations, iterations ) # Return execution time in microseconds - exec_time = testing.benchmark( + return testing.benchmark( compiled_fn, workspace_generator=generate_tensors, workspace_count=workspace_count, @@ -1579,8 +1669,6 @@ def run( warmup_iterations=warmup_iterations, iterations=iterations, ) - print(f"[DSL INFO] Execution time: {exec_time} microseconds per iteration") - return exec_time def _parse_comma_separated_ints(s: str) -> Tuple[int, ...]: @@ -1654,6 +1742,12 @@ def prepare_parser(): default=False, help="Use circular buffer tensor sets to ensure L2 cold cache", ) + + return parser + + +if __name__ == "__main__": + parser = prepare_parser() parser.add_argument( "--mma_tiler_mn", type=_parse_comma_separated_ints, @@ -1661,40 +1755,6 @@ def prepare_parser(): help="Mma tile shape (comma-separated)", ) - return parser - - -def create_tensors(l, m, n, k, a_major, b_major, c_major, ab_dtype, c_dtype): - torch.manual_seed(1111) - - a_torch_cpu = cutlass_torch.matrix(l, m, k, a_major == "m", ab_dtype) - b_torch_cpu = cutlass_torch.matrix(l, n, k, b_major == "n", ab_dtype) - c_torch_cpu = cutlass_torch.matrix(l, m, n, c_major == "m", c_dtype) - - a_tensor, _ = cutlass_torch.cute_tensor_like( - a_torch_cpu, ab_dtype, is_dynamic_layout=True, assumed_align=16 - ) - b_tensor, _ = cutlass_torch.cute_tensor_like( - b_torch_cpu, ab_dtype, is_dynamic_layout=True, assumed_align=16 - ) - c_tensor, c_torch_gpu = cutlass_torch.cute_tensor_like( - c_torch_cpu, c_dtype, is_dynamic_layout=True, assumed_align=16 - ) - - return ( - a_tensor, - b_tensor, - c_tensor, - a_torch_cpu, - b_torch_cpu, - c_torch_cpu, - c_torch_gpu, - ) - - -if __name__ == "__main__": - parser = prepare_parser() - args = parser.parse_args() if len(args.mnkl) != 4: @@ -1739,4 +1799,5 @@ if __name__ == "__main__": args.use_cold_l2, args.benchmark == "default", ) - print(f"Execution time: {exec_time} microseconds per iteration") + print(f"Execution time: {exec_time} seconds") + print("PASS") diff --git a/examples/python/CuTeDSL/experimental/blackwell/dense_gemm_ptr_array.py b/examples/python/CuTeDSL/experimental/blackwell/dense_gemm_ptr_array.py index 2ec3d73a..bd4e09c1 100755 --- a/examples/python/CuTeDSL/experimental/blackwell/dense_gemm_ptr_array.py +++ b/examples/python/CuTeDSL/experimental/blackwell/dense_gemm_ptr_array.py @@ -198,10 +198,7 @@ class DenseGemmPtrArrayKernel: ) # UMMA ACC TMEM Layout - acc_shape = tiled_mma.partition_shape_C(mnk_tiler[:2]) - tmem_layout = tiled_mma.make_fragment_C( - cute.append(acc_shape, acc_stage) - ).layout + tmem_layout = cute_ext.make_tmem_layout_acc(tiled_mma, mnk_tiler, acc_stage) # Allocate UMMA Buffers bufferA = cute_ext.allocate( @@ -251,11 +248,9 @@ class DenseGemmPtrArrayKernel: tiled_copy_t2r = cute.nvgpu.tcgen05.make_tmem_copy(copy_atom_t2r, acc_epi_div) # Calculate the per thread destination size per iteration for output of TMEM and input of SMEM - thr_copy_t2r = tiled_copy_t2r.get_slice(tid_x) gC_mnl_epi = cute.flat_divide(gD_tile, epi_tile) - tTR_gC = thr_copy_t2r.partition_D(gC_mnl_epi) - acc_d_rmem_layout = cute.make_fragment_like( - tTR_gC[(None, None, None, 0, 0)].layout + acc_d_rmem_layout = cute_ext.make_t2r_rmem_layout( + tiled_copy_t2r, gC_mnl_epi, tid_x ) # Allocate RMEM buffers diff --git a/examples/python/CuTeDSL/helpers/sparse_utils.py b/examples/python/CuTeDSL/helpers/sparse_utils.py deleted file mode 100644 index 24b3f791..00000000 --- a/examples/python/CuTeDSL/helpers/sparse_utils.py +++ /dev/null @@ -1,457 +0,0 @@ -import numpy as np -import cutlass -import cutlass.cute as cute -from cutlass.cute.runtime import from_dlpack -import torch - - -@cute.jit -def print_tensor_dlpack(src: cute.Tensor): - print(src) - cute.print_tensor(src) - - -# Sparse emulation -class SparseEmulation: - def __init__(self, M: int, N: int, K: int, L: int): - self.M = M - self.N = N - self.K = K - self.L = L - - @cute.jit - def __call__(self, a: cute.Tensor, b: cute.Tensor, d: cute.Tensor, e: cute.Tensor): - """Sparse emulation""" - num_threads = 128 - grid = (cute.ceil_div(self.M, num_threads), 1, 1) - block = (num_threads, 1, 1) - self.kernel(a, b, d, e).launch(grid=grid, block=block) - return - - @cute.kernel - def kernel(self, a: cute.Tensor, b: cute.Tensor, d: cute.Tensor, e: cute.Tensor): - """CUDA kernel to emulate sparse tensor core""" - tidx, tidy, tidz = cute.arch.thread_idx() - bidx, bidy, bidz = cute.arch.block_idx() - - row_idx = tidx + bidx * self.M - meta_idx = self.K // 4 // 8 - if row_idx < self.M: - # each thread process 1 row - for col in range(self.N): - # each meta_idx stands for 32 elements - for e_idx in range(meta_idx): - meta_val = e[(row_idx, e_idx)] - for k in range(8): - # each k stands for 4 elements - meta_row = (meta_val >> (k * 4)) & 0xF - idx0 = meta_row & 0x3 - idx1 = (meta_row >> 2) & 0x3 - # calculate the idx in b tensor which has value in A tensor - km = e_idx * 16 + k * 2 - km_1 = km + 1 - kn = e_idx * 32 + k * 4 + idx0 - kn_1 = e_idx * 32 + k * 4 + idx1 - d[row_idx, col] += a[row_idx, km] * b[col, kn] - d[row_idx, col] += a[row_idx, km_1] * b[col, kn_1] - return - - -# Compressor -# compress a sparse tensor to a dense tensor && generate metadata -class Compressor: - def __init__(self, M: int, K: int, L: int): - self.M = M - self.K = K - self.L = L - self.pos_map = { - 0x4: [0, 1], - 0x8: [0, 2], - 0xC: [0, 3], - 0x9: [1, 2], - 0xD: [1, 3], - 0xE: [2, 3], - } - - @cute.jit - def _init__(self, a: cute.Tensor): - self.__init__(a.shape[0], a.shape[1], a.shape[2]) - - def compress(self, a, a_compressed, meta, run_on_cpu: bool): - if run_on_cpu: - if a.device.type != "cpu": - raise ValueError("a must be on cpu") - return self.__compress_on_cpu(a, a_compressed, meta) - else: - if a.device.type != "cuda": - raise ValueError("a must be on cuda") - return self.__compress_on_cuda(a, a_compressed, meta) - - def __compress_on_cpu(self, a, a_compressed, meta): - """ - compress the tensor on cpu - # Convert to 4-bit metadata value - # The metadata value represents which 2 elements are non-zero - # 0x4: [1,1,0,0] - first two elements are non-zero - # 0x8: [1,0,1,0] - first and third elements are non-zero - # 0xC: [1,0,0,1] - first and fourth elements are non-zero - # 0x9: [0,1,1,0] - second and third elements are non-zero - # 0xD: [0,1,0,1] - second and fourth elements are non-zero - # 0xE: [0,0,1,1] - third and fourth elements are non-zero - # special case: - # [0,0,0,0] == [0,0,1,1] - # [1,0,0,0] == [1,0,0,1] - # [0,1,0,0] == [0,1,0,1] - # [0,0,1,0] == [0,0,1,1] - # [0,0,0,1] == [0,0,1,1] - """ - M, K = a.shape - assert a_compressed.shape == ( - M, - K // 2, - ), f"Expected a_compressed shape {(M, K // 2)}, got {a_compressed.shape}" - assert meta.shape == ( - M, - K // 4 // 8, - ), f"Expected meta shape {(M, K // 4 // 8)}, got {meta.shape}" - for m in range(M): - k_meta = 0 - for k in range(0, K, 4): - chunk = a[m, k : k + 4] - - non_zero_indices = torch.nonzero(chunk).squeeze() - meta_val = 0xE - if torch.equal(non_zero_indices, torch.tensor([0, 1])): - meta_val = 0x4 - elif torch.equal(non_zero_indices, torch.tensor([0, 2])): - meta_val = 0x8 - elif torch.equal(non_zero_indices, torch.tensor([0, 3])) or torch.equal( - non_zero_indices, torch.tensor(0) - ): - meta_val = 0xC - elif torch.equal(non_zero_indices, torch.tensor([1, 2])): - meta_val = 0x9 - elif torch.equal(non_zero_indices, torch.tensor([1, 3])) or torch.equal( - non_zero_indices, torch.tensor(1) - ): - meta_val = 0xD - elif torch.equal(non_zero_indices, torch.tensor([2, 3])) or torch.equal( - non_zero_indices, torch.tensor(2) - ): - meta_val = 0xE - elif torch.equal(non_zero_indices, torch.tensor([])) or torch.equal( - non_zero_indices, torch.tensor(3) - ): - meta_val = 0xE - else: - raise ValueError(f"Invalid non-zero pattern: {non_zero_indices}") - meta_idx = k // 4 // 8 - meta_bit_pos = (k // 4) % 8 - if k_meta == meta_idx: - k_meta = meta_idx + 1 - meta[m, meta_idx] = 0 - meta[m, meta_idx] |= meta_val << (meta_bit_pos * 4) - compressed_idx = k // 2 - index = self.pos_map[meta_val] - a_compressed[m, compressed_idx] = chunk[index[0]] - a_compressed[m, compressed_idx + 1] = chunk[index[1]] - - def __compress_on_cuda(self, a, a_compressed, meta): - """ - compress the tensor on cuda - """ - a_tensor = from_dlpack(a) - a_compressed_tensor = from_dlpack(a_compressed) - meta_tensor = from_dlpack(meta) - self.compress_on_cuda_impl(a_tensor, a_compressed_tensor, meta_tensor) - return - - @cute.jit - def compress_on_cuda_impl( - self, a: cute.Tensor, a_compressed: cute.Tensor, meta: cute.Tensor - ): - """Compress the input tensor using the metadata""" - num_threads = 128 - grid = (cute.ceil_div(self.M, num_threads), 1, 1) - block = (num_threads, 1, 1) - self.compressor_impl(a, a_compressed, meta).launch(grid=grid, block=block) - - @cute.kernel - def compressor_impl( - self, a: cute.Tensor, a_compressed: cute.Tensor, meta: cute.Tensor - ): - """CUDA kernel to compress the tensor""" - tidx, tidy, tidz = cute.arch.thread_idx() - bidx, bidy, bidz = cute.arch.block_idx() - m = a.shape[0] - k = a.shape[1] - - # each thread process 1 row - row_idx = tidx + bidx * self.M - meta_idx = self.K // 4 // 8 - if row_idx < self.M: - # each meta_idx stands for 32 elements - for i in range(meta_idx): - meta[row_idx, i] = 0 - # each k stands for 4 elements - for j in range(8): - val = a[row_idx, i * 32 + j * 4] - val_1 = a[row_idx, i * 32 + j * 4 + 1] - val_2 = a[row_idx, i * 32 + j * 4 + 2] - val_3 = a[row_idx, i * 32 + j * 4 + 3] - value_idx = 0 - value_idx_1 = 0 - value_idx_2 = 0 - value_idx_3 = 0 - pos0 = 0 - pos1 = 0 - if val != 0: - value_idx = 1 - pos0 = 0 - if val_1 != 0: - value_idx_1 = 1 - if val_2 != 0: - value_idx_2 = 1 - if val_3 != 0: - value_idx_3 = 1 - pos = [value_idx, value_idx_1, value_idx_2, value_idx_3] - tmp = 0 - if pos == [0, 0, 0, 0]: - tmp = 0xE - pos0 = 2 - pos1 = 3 - elif pos == [1, 0, 0, 0]: - tmp = 0xC - pos0 = 0 - pos1 = 3 - elif pos == [0, 1, 0, 0]: - tmp = 0xD - pos0 = 1 - pos1 = 3 - elif pos == [0, 0, 1, 0]: - tmp = 0xE - pos0 = 2 - pos1 = 3 - elif pos == [0, 0, 0, 1]: - tmp = 0xE - pos0 = 2 - pos1 = 3 - elif pos == [1, 1, 0, 0]: - tmp = 0x4 - pos0 = 0 - pos1 = 1 - elif pos == [1, 0, 1, 0]: - tmp = 0x8 - pos0 = 0 - pos1 = 2 - elif pos == [1, 0, 0, 1]: - tmp = 0xC - pos0 = 0 - pos1 = 3 - elif pos == [0, 1, 1, 0]: - tmp = 0x9 - pos0 = 1 - pos1 = 2 - elif pos == [0, 1, 0, 1]: - tmp = 0xD - pos0 = 1 - pos1 = 3 - elif pos == [0, 0, 1, 1]: - tmp = 0xE - pos0 = 2 - pos1 = 3 - # cute.printf(row_idx, cutlass.Float32(val), cutlass.Float32(val_1), cutlass.Float32(val_2), cutlass.Float32(val_3), tmp) - meta[row_idx, i] |= tmp << (j * 4) - - a_compressed[row_idx, i * 16 + j * 2] = a[ - row_idx, i * 32 + j * 4 + pos0 - ] - a_compressed[row_idx, i * 16 + j * 2 + 1] = a[ - row_idx, i * 32 + j * 4 + pos1 - ] - - return - - -# SparseUtils is used to generate sparse tensor -# format torch.Tensor -class SparseUtils: - #!brief: SparseUtils is used to generate sparse tensor - #!param: M: int, K: int, L: int, dtype: cutlass.DataType - def __init__(self, M: int, K: int, L: int, dtype): - self.M = M - self.K = K - self.L = L - self.dtype = dtype - self.meta_data = self._generate_meta_data_4_2() - self._use_specific_meta_data = False - - #!brief: cast cutlass.DataType to torch.Tensor - def _get_type(self): - if self.dtype == cutlass.Float16: - return torch.float16 - elif self.dtype == cutlass.Float32: - return torch.float32 - elif self.dtype == cutlass.Int8: - return torch.int8 - else: - raise ValueError(f"Unsupported dtype: {self.dtype}") - - def _generate_meta_data_4_2(self): - # metadata for 4:2 sparse will in range( 4,8,9,c,d,e) - # represents - # 0: [1,1,0,0] no zero pos 00,01 -> 0100 = 4 - # 1: [1,0,1,0] no zero pos 00,10 -> 1000 = 8 - # 2: [1,0,0,1] no zero pos 00,11 -> 1100 = c - # 3: [0,1,1,0] no zero pos 01,10 -> 1001 = 9 - # 4: [0,1,0,1] no zero pos 01,11 -> 1101 = d - # 5: [0,0,1,1] no zero pos 10,11 -> 1011 = e - meta_value = [0x4, 0x8, 0x9, 0xC, 0xD, 0xE] - # 4:2 sparse, so each chunk is 4 elements, map to 4 bits - K_NumChunk = self.K // 4 - meta_data = np.random.choice( - meta_value, size=(self.M, K_NumChunk), replace=True - ) - meta_data = torch.from_numpy( - np.array(meta_data).astype(np.uint8).reshape(self.M, K_NumChunk) - ) - return meta_data - - #!brief: pack meta data - def _pack_meta_data(self): - tmp = [] - K_NumChunk = self.K // 4 - for i in range(self.M): - for j in range(K_NumChunk // 8): - v = 0 - for k in range(8): - vv = int(self.meta_data[i, j * 8 + k] & 0xF) - tt = vv << (k * 4) - v = v | tt - tmp.append(v) - # debug print - # print([hex(vt) for vt in tmp]) - result = torch.from_numpy( - np.array(tmp).astype(np.uint32).reshape(self.M, K_NumChunk // 8) - ) - return result - - #!brief: use specific meta data - def use_specific_meta_data(self, meta_data: torch.Tensor = None): - if meta_data is not None: - self.meta_data = meta_data - self._use_specific_meta_data = True - - #!brief: generate sparse tensor with tensor - #!param: a: torch.Tensor - #!param: run_on_cpu: bool - #!return: torch.Tensor - def generate_sparse_4_2_tensor_with_tensor(self, a, run_on_cpu): - if run_on_cpu: - if a.device.type != "cpu": - raise ValueError("a must be on cpu") - return self.__generate_sparse_tensor_cpu(a) - else: - if a.device.type != "cuda": - raise ValueError("a must be on cuda") - a_tensor = from_dlpack(a) - packed_meta_data = self._pack_meta_data() - meta_tensor = from_dlpack(packed_meta_data.cuda()) - self.__generate_sparse_tensor_cuda(a_tensor, meta_tensor) - return a - - #!brief: generate sparse tensor - #!param: run_on_cpu: bool - #!return: torch.Tensor - def generate_4_2_sparse_tensor(self, run_on_cpu): - dtype = self._get_type() - a = torch.empty(self.M, self.K).random_(-5, 5).to(dtype) - if run_on_cpu: - return self.generate_sparse_4_2_tensor_with_tensor(a, run_on_cpu) - else: - return self.generate_sparse_4_2_tensor_with_tensor(a.cuda(), run_on_cpu) - - #!brief: generate sparse tensor on cpu - #!param: a: torch.Tensor - #!return: torch.Tensor - def __generate_sparse_tensor_cpu(self, a): - if not self._use_specific_meta_data: - for m in range(self.M): - for k in range(0, self.K, 4): - # random choose 2 zero positions - zero_indices = torch.randperm(4)[:2] - a[m, k + zero_indices[0]] = 0 - a[m, k + zero_indices[1]] = 0 - return a - else: - # use specific meta data - tensor_mask = [] - for i in range(self.M): - for j in range(self.K // 4): - meta_val = self.meta_data[i, j] - tmp = [] - if meta_val == 0x4: - tmp = [1, 1, 0, 0] - elif meta_val == 0x8: - tmp = [1, 0, 1, 0] - elif meta_val == 0xC: - tmp = [1, 0, 0, 1] - elif meta_val == 0x9: - tmp = [0, 1, 1, 0] - elif meta_val == 0xD: - tmp = [0, 1, 0, 1] - elif meta_val == 0xE: - tmp = [0, 0, 1, 1] - tensor_mask.extend(tmp) - a = torch.reshape(a, (-1,)) - mask = torch.tensor(tensor_mask) - a = a * mask - a = torch.reshape(a, (self.M, self.K)) - return a - - @cute.jit - def __generate_sparse_tensor_cuda(self, a: cute.Tensor, meta: cute.Tensor): - """Generate a sparse tensor from a dense tensor using metadata""" - assert a.shape[0] == self.M and a.shape[1] == self.K - assert meta.shape[0] == self.M and meta.shape[1] == self.K // 4 // 8 - num_threads = 128 - grid = (cute.ceil_div(self.M, num_threads), 1, 1) - block = (num_threads, 1, 1) - self.kernel(a, meta).launch(grid=grid, block=block) - - @cute.kernel - def kernel(self, a: cute.Tensor, meta: cute.Tensor): - """Apply sparsity mask to input tensor using metadata""" - tidx, tidy, tidz = cute.arch.thread_idx() - bidx, bidy, bidz = cute.arch.block_idx() - - # each thread process 1 ro - row_idx = tidx + bidx * self.M - meta_idx = self.K // 4 // 8 - # each thread process 1 row - if row_idx < self.M: - # iterate over each chunk(32 elements) - for i in range(meta_idx): - meta_val = meta[(row_idx, i)] - # iterate over each sparse pattern(4 elements) - for j in range(8): - meta_row = (meta_val >> (j * 4)) & 0xF - idx0 = meta_row & 0x3 - idx1 = (meta_row >> 2) & 0x3 - r_id0 = 0 - r_id1 = 0 - # r_id is the idx that value is 0 - if idx0 >= 2 and idx1 >= 2: - r_id0 = 0 - r_id1 = 1 - elif idx0 <= 1 and idx1 <= 1: - r_id0 = 2 - r_id1 = 3 - else: - r_id0 = idx0 ^ 0b1 - r_id1 = idx1 ^ 0b1 - row_id0 = r_id0 + i * 32 + j * 4 - row_id1 = r_id1 + i * 32 + j * 4 - a[row_idx, row_id0] = self.dtype(0.0) - a[row_idx, row_id1] = self.dtype(0.0) - return diff --git a/examples/python/CuTeDSL/helpers/test_sparse_utils.py b/examples/python/CuTeDSL/helpers/test_sparse_utils.py deleted file mode 100644 index 3264f191..00000000 --- a/examples/python/CuTeDSL/helpers/test_sparse_utils.py +++ /dev/null @@ -1,104 +0,0 @@ -import sparse_utils as su -import cutlass -import torch -from cutlass.cute.runtime import from_dlpack -import numpy as np -import pytest - - -@pytest.mark.L0 -def test_sparse_cpu(): - M = 128 - N = 32 - K = 32 - L = 1 - debug = False - # generate sparse tensor - a = torch.empty(M, K).random_(-5, 5).to(torch.float16) - sparse_utils = su.SparseUtils(M, K, L, cutlass.Float16) - if debug: - sparse_utils.use_specific_meta_data() - a_gen_from_cpu = sparse_utils.generate_sparse_4_2_tensor_with_tensor(a, True) - # print(a_gen_from_cpu) - # generate compressed tensor and meta data - a_compressed_cpu = torch.empty(M, K // 2).to(torch.float16) - meta_data_cpu = torch.empty(M, K // 4 // 8).to(torch.uint32) - compressor = su.Compressor(M, K, L) - compressor.compress(a_gen_from_cpu, a_compressed_cpu, meta_data_cpu, True) - # # test with gemm - b = torch.empty(N, K).random_(-5, 5).to(torch.float16).cuda() - d = torch.empty(M, N).zero_().to(torch.float16).cuda() - b_tensor = from_dlpack(b) - d_tensor = from_dlpack(d) - a_compressed_cpu_tensor = from_dlpack(a_compressed_cpu.cuda()) - meta_data_cpu_tensor = from_dlpack(meta_data_cpu.cuda()) - sparse_emulation = su.SparseEmulation(M, N, K, 1) - sparse_emulation(a_compressed_cpu_tensor, b_tensor, d_tensor, meta_data_cpu_tensor) - - ref = torch.einsum("mk,nk->mn", a_gen_from_cpu.cpu(), b.cpu()) - if debug: - a_ori = a_gen_from_cpu.cpu().numpy() - np.savetxt("a.txt", a_ori, fmt="%f") - a_compressed_cpu_ori = a_compressed_cpu.cpu().numpy() - np.savetxt("a_compressed_cpu.txt", a_compressed_cpu_ori, fmt="%f") - meta_data_cpu_ori = meta_data_cpu.cpu().numpy() - np.savetxt("meta_data_cpu.txt", meta_data_cpu_ori, fmt="%f") - d_ori = d.cpu().numpy() - np.savetxt("d.txt", d_ori, fmt="%f") - ref_ori = ref.cpu().numpy() - np.savetxt("ref.txt", ref_ori, fmt="%f") - torch.testing.assert_close(d.cpu(), ref) - print("cpu d == ref") - - -@pytest.mark.L0 -def test_sparse_cuda(): - M = 128 - N = 32 - K = 32 - L = 1 - debug = False - sparse_utils = su.SparseUtils(M, K, L, cutlass.Float16) - if debug: - sparse_utils.use_specific_meta_data() - # generate sparse tensor - a = torch.empty(M, K).random_(-5, 5).to(torch.float16).cuda() - a_gen_from_cuda = sparse_utils.generate_4_2_sparse_tensor(False) - # print(a_gen_from_cuda) - # generate compressed tensor and meta data - a_compressed_cuda = torch.empty(M, K // 2).to(torch.float16).cuda() - meta_data_cuda = torch.empty(M, K // 4 // 8).to(torch.uint32).cuda() - compressor = su.Compressor(M, K, L) - compressor.compress(a_gen_from_cuda, a_compressed_cuda, meta_data_cuda, False) - # test with gemm - b = torch.empty(N, K).random_(-5, 5).to(torch.float16).cuda() - d = torch.empty(M, N).zero_().to(torch.float16).cuda() - b_tensor = from_dlpack(b) - d_tensor = from_dlpack(d) - a_compressed_cuda_tensor = from_dlpack(a_compressed_cuda) - meta_data_cuda_tensor = from_dlpack(meta_data_cuda) - sparse_emulation = su.SparseEmulation(M, N, K, 1) - sparse_emulation( - a_compressed_cuda_tensor, b_tensor, d_tensor, meta_data_cuda_tensor - ) - - ref = torch.einsum("mk,nk->mn", a_gen_from_cuda.cpu(), b.cpu()) - if debug: - a_ori = a_gen_from_cuda.cpu().numpy() - np.savetxt("a.txt", a_ori, fmt="%f") - a_compressed_cuda_ori = a_compressed_cuda.cpu().numpy() - np.savetxt("a_compressed_cuda.txt", a_compressed_cuda_ori, fmt="%f") - meta_data_cuda_ori = meta_data_cuda.cpu().numpy() - np.savetxt("meta_data_cuda.txt", meta_data_cuda_ori, fmt="%f") - d_ori = d.cpu().numpy() - np.savetxt("d.txt", d_ori, fmt="%f") - ref_ori = ref.cpu().numpy() - np.savetxt("ref.txt", ref_ori, fmt="%f") - torch.testing.assert_close(d.cpu(), ref) - print("cuda d == ref") - - -if __name__ == "__main__": - cutlass.cuda.initialize_cuda_context() - test_sparse_cpu() - test_sparse_cuda() diff --git a/examples/python/CuTeDSL/hopper/dense_gemm.py b/examples/python/CuTeDSL/hopper/dense_gemm.py index c00ead6f..96ad8338 100644 --- a/examples/python/CuTeDSL/hopper/dense_gemm.py +++ b/examples/python/CuTeDSL/hopper/dense_gemm.py @@ -31,15 +31,12 @@ from typing import Tuple, Type import math import cuda.bindings.driver as cuda -import torch - import cutlass import cutlass.cute as cute import cutlass.cute.testing as testing import cutlass.utils as utils import cutlass.pipeline as pipeline from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait -import cutlass.torch as cutlass_torch from cutlass.cute.runtime import from_dlpack import cutlass.utils.hopper_helpers as sm90_utils @@ -1006,7 +1003,10 @@ class HopperWgmmaGemmKernel: tiled_copy_r2s, tRS_rD_out, tRS_sD[(None, None, None, epi_buffer)] ) - cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) # barrier for sync pipeline.sync(barrier_id=1) @@ -1431,6 +1431,9 @@ def run( :rtype: float """ + import torch + import cutlass.torch as cutlass_torch + print("Running Hopper Dense GEMM with:") print(f"mnkl: {mnkl}") print( @@ -1519,7 +1522,7 @@ def run( gemm = HopperWgmmaGemmKernel(acc_dtype, tile_shape_mn, cluster_shape_mn) - torch_stream = torch.cuda.Stream() + torch_stream = torch.cuda.current_stream() stream = cuda.CUstream(torch_stream.cuda_stream) # compile gemm kernel compiled_gemm = cute.compile(gemm, mA, mB, mC, stream) diff --git a/examples/python/CuTeDSL/hopper/dense_gemm_persistent.py b/examples/python/CuTeDSL/hopper/dense_gemm_persistent.py index 5b999c0b..a8119c8a 100644 --- a/examples/python/CuTeDSL/hopper/dense_gemm_persistent.py +++ b/examples/python/CuTeDSL/hopper/dense_gemm_persistent.py @@ -31,14 +31,11 @@ from typing import Optional, Tuple, Type import math import cuda.bindings.driver as cuda -import torch - import cutlass import cutlass.cute as cute import cutlass.cute.testing as testing import cutlass.pipeline as pipeline from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait -import cutlass.torch as cutlass_torch import cutlass.utils as utils import cutlass.utils.hopper_helpers as sm90_utils @@ -952,7 +949,10 @@ class HopperWgmmaGemmPersistentKernel: tRS_sD[(None, None, None, epi_buffer)], ) - cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) self.epilog_sync_barrier.arrive_and_wait() gmem_coord = epi_tile_layout.get_hier_coord(epi_idx) @@ -1465,6 +1465,8 @@ def run( :return: Execution time of the GEMM kernel in microseconds :rtype: float """ + import torch + import cutlass.torch as cutlass_torch print("Running Hopper Persistent Dense GEMM with:") print(f"mnkl: {mnkl}") diff --git a/examples/python/CuTeDSL/hopper/fmha.py b/examples/python/CuTeDSL/hopper/fmha.py index 9efce75b..e384c53b 100644 --- a/examples/python/CuTeDSL/hopper/fmha.py +++ b/examples/python/CuTeDSL/hopper/fmha.py @@ -83,9 +83,6 @@ import sys import time from typing import Type, Tuple, Optional -import torch - - import cuda.bindings.driver as cuda import cutlass @@ -96,7 +93,6 @@ import cutlass.cute.nvgpu.warpgroup as warpgroup import cutlass.utils as utils import cutlass.pipeline as pipeline from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait -import cutlass.torch as cutlass_torch from cutlass._mlir.dialects import math as _math import cutlass.utils.hopper_helpers as sm90_utils @@ -598,6 +594,7 @@ class HopperFusedMultiHeadAttentionForward: k_smem_layout_staged.outer, swizzle=k_smem_layout_staged.inner ) # (MMA, MMA_K, MMA_D, PIPE) + # Adjust swizzle info to reuse smem sV_ptr = cute.recast_ptr(sK.iterator, v_smem_layout_staged.inner) sV = cute.make_tensor(sV_ptr, v_smem_layout_staged.outer) @@ -648,11 +645,19 @@ class HopperFusedMultiHeadAttentionForward: producer_warp_role = warp_idx % 4 # self.num_warps_per_warp_group + # Fence the mbarrier init to ensure all mbarrier initializations are visible + # to all threads. This is critical for FP8 performance - without this fence, + # the compiler may generate software polling loops instead of hardware waits. + cute.arch.mbarrier_init_fence() + # 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 - pipeline_init_arrive(cluster_shape_mn=self.cluster_shape_mnk, is_relaxed=True) - pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mnk) + if cute.size(self.cluster_shape_mnk) > 1: + cute.arch.cluster_arrive_relaxed() + cute.arch.cluster_wait() + else: + cute.arch.sync_threads() if warp_idx == 0: cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_q) @@ -1164,7 +1169,10 @@ class HopperFusedMultiHeadAttentionForward: tRS_sD[(None, None, None, epi_buffer, warp_group_idx - 1)], ) - cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) pipeline.arrive_and_wait( barrier_id=warp_group_idx, num_threads=self.num_threads_per_warp_group, @@ -1935,6 +1943,9 @@ def run( :return: Execution time of the FMHA kernel in microseconds :rtype: float """ + import torch + import cutlass.torch as cutlass_torch + print("Running Hopper SM90 FMHA test with:") print(f" q_shape: {q_shape}") print(f" k_shape: {k_shape}") diff --git a/examples/python/CuTeDSL/jax/cutlass_call_sharding.py b/examples/python/CuTeDSL/jax/cutlass_call_sharding.py index a2d03c34..e40687ff 100644 --- a/examples/python/CuTeDSL/jax/cutlass_call_sharding.py +++ b/examples/python/CuTeDSL/jax/cutlass_call_sharding.py @@ -31,7 +31,7 @@ import argparse import jax import jax.numpy as jnp -from jax.sharding import Mesh, NamedSharding, PartitionSpec as P +from jax.sharding import NamedSharding, PartitionSpec as P, AxisType from jax.experimental.custom_partitioning import custom_partitioning import cutlass @@ -42,8 +42,8 @@ import cuda.bindings.driver as cuda """ -Examples of combining jax.jit and jax.shard_map for sharding and executing kernels -across multiple GPU devices. +Examples of combining jax.jit, jax.shard_map and custom_partitioning for sharding +and executing kernels across multiple GPU devices. To run this example: @@ -84,25 +84,57 @@ def launch( ) +def sharded_cutlass_call_impl(a_block, b_block): + """The sharded implementation that operates on a single device.""" + call = cjax.cutlass_call( + launch, + use_static_tensors=True, + output_shape_dtype=jax.ShapeDtypeStruct(a_block.shape, a_block.dtype), + ) + ref_result = a_block + b_block + return call(a_block, b_block), ref_result + + +@custom_partitioning +def custom_shared_call(a, b): + return sharded_cutlass_call_impl(a, b) + + +def custom_shared_call_partitioner(mesh, arg_shapes, result_shape): + arg_shardings = jax.tree.map(lambda x: x.sharding, arg_shapes) + result_shardings = tuple([arg_shardings[0]] * len(result_shape)) + + def lower_fn(*args): + return sharded_cutlass_call_impl(*args) + + return mesh, lower_fn, result_shardings, arg_shardings + + +custom_shared_call.def_partition(custom_shared_call_partitioner) + + def run_example(): # Create a device mesh with one axis b ngpu = jax.device_count() - mesh = jax.make_mesh((ngpu,), "b") + mesh = jax.make_mesh((ngpu,), "b", axis_types=(AxisType.Explicit,)) if ngpu == 1: print("Note: only 1 GPU was detected.") # We will shard our 3D tensors over b sharding = P("b", None, None) + named_sharding = NamedSharding(mesh, sharding) - @partial(jax.jit, static_argnums=[0, 1]) + print("Testing shard_map...") + + @partial( + jax.jit, static_argnums=[0, 1], out_shardings=(named_sharding, named_sharding) + ) def allocate_sharded_tensors(shape, dtype): key = jax.random.key(1123) - a_key, b_keys = jax.random.split(key, 2) + a_key, b_key = jax.random.split(key, 2) a = create_tensor(shape, dtype, a_key) - b = create_tensor(shape, dtype, b_keys) - a = jax.lax.with_sharding_constraint(a, NamedSharding(mesh, sharding)) - b = jax.lax.with_sharding_constraint(b, NamedSharding(mesh, sharding)) + b = create_tensor(shape, dtype, b_key) return a, b @jax.jit @@ -115,13 +147,7 @@ def run_example(): out_specs=(sharding, sharding), ) def sharded_call(a_block, b_block): - call = cjax.cutlass_call( - launch, - use_static_tensors=True, - output_shape_dtype=jax.ShapeDtypeStruct(a_block.shape, a_block.dtype), - ) - ref_result = a_block + b_block - return call(a_block, b_block), ref_result + return sharded_cutlass_call_impl(a_block, b_block) return sharded_call(a, b) @@ -134,6 +160,17 @@ def run_example(): assert jnp.allclose(c, c_ref) + print("Testing custom_partitioning...") + + # Test custom_partitioning implementation which should produce identical results + @jax.jit + def compute_cp(a, b): + return custom_shared_call(a, b) + + c, c_ref = compute_cp(a, b) + + assert jnp.allclose(c, c_ref) + if __name__ == "__main__": run_example() diff --git a/examples/python/CuTeDSL/notebooks/hello_world.ipynb b/examples/python/CuTeDSL/notebooks/hello_world.ipynb index 218378ad..6bf35b76 100644 --- a/examples/python/CuTeDSL/notebooks/hello_world.ipynb +++ b/examples/python/CuTeDSL/notebooks/hello_world.ipynb @@ -55,7 +55,7 @@ " # Get the x component of the thread index (y and z components are unused)\n", " tidx, _, _ = cute.arch.thread_idx()\n", " # Only the first thread (thread 0) prints the message\n", - " if tidx == 0:\n", + " if cutlass.dynamic_expr(tidx == 0):\n", " cute.printf(\"Hello world\")" ] }, @@ -142,8 +142,6 @@ "from cutlass.cute import KeepPTX, KeepCUBIN\n", "\n", "print(\"Compiling with PTX/CUBIN dumped...\")\n", - "# Alternatively, compile with string based options like\n", - "# cute.compile(hello_world, options=\"--keep-ptx --keep-cubin\") would also work.\n", "hello_world_compiled_ptx_on = cute.compile[KeepPTX, KeepCUBIN](hello_world)\n", "\n", "# Run the pre-compiled version\n", diff --git a/include/cutlass/arch/grid_dependency_control.h b/include/cutlass/arch/grid_dependency_control.h index 912b635a..63a7714e 100644 --- a/include/cutlass/arch/grid_dependency_control.h +++ b/include/cutlass/arch/grid_dependency_control.h @@ -62,6 +62,8 @@ (defined(__CUDA_ARCH_FEAT_SM100_ALL) || CUDA_ARCH_FAMILY(1000))) || \ (__CUDA_ARCH__ == 1010 &&\ (defined(__CUDA_ARCH_FEAT_SM101_ALL) || CUDA_ARCH_FAMILY(1010))) || \ + (__CUDA_ARCH__ == 1100 &&\ + (defined(__CUDA_ARCH_FEAT_SM110_ALL) || CUDA_ARCH_FAMILY(1100))) || \ (__CUDA_ARCH__ == 1030 &&\ (defined(__CUDA_ARCH_FEAT_SM103_ALL) || CUDA_ARCH_FAMILY(1030))) || \ (__CUDA_ARCH__ == 1200 &&\ diff --git a/include/cutlass/conv/kernel/sm100_implicit_gemm_tma_warpspecialized.hpp b/include/cutlass/conv/kernel/sm100_implicit_gemm_tma_warpspecialized.hpp index fbf88553..494ffe7a 100644 --- a/include/cutlass/conv/kernel/sm100_implicit_gemm_tma_warpspecialized.hpp +++ b/include/cutlass/conv/kernel/sm100_implicit_gemm_tma_warpspecialized.hpp @@ -276,7 +276,11 @@ public: static constexpr int MaxClusterSize = 16; implementable &= size(args.hw_info.cluster_shape) <= MaxClusterSize; implementable &= size(args.hw_info.cluster_shape_fallback) <= MaxClusterSize; - implementable &= cutlass::detail::preferred_cluster_can_implement(args.hw_info.cluster_shape, args.hw_info.cluster_shape_fallback); + // Early return if cluster shape validation failed to avoid division by zero below + if (not cutlass::detail::preferred_cluster_can_implement(args.hw_info.cluster_shape, args.hw_info.cluster_shape_fallback)) { + CUTLASS_TRACE_HOST(" CAN IMPLEMENT: Invalid dynamic cluster shape\n"); + return false; + } } auto cluster_shape = cutlass::detail::select_cluster_shape(ClusterShape{}, args.hw_info.cluster_shape); diff --git a/include/cutlass/gemm/collective/sm100_mma_array_warpspecialized_emulated.hpp b/include/cutlass/gemm/collective/sm100_mma_array_warpspecialized_emulated.hpp index bdeae396..15dd91bc 100644 --- a/include/cutlass/gemm/collective/sm100_mma_array_warpspecialized_emulated.hpp +++ b/include/cutlass/gemm/collective/sm100_mma_array_warpspecialized_emulated.hpp @@ -181,6 +181,11 @@ struct CollectiveMma< static constexpr uint32_t NumTransformationThreads = 128; static constexpr uint32_t NumAccumThreads = 128; + // Register reconfiguration + static constexpr uint32_t GenericRegisterRequirement = 64; + static constexpr uint32_t TransformRegisterRequirement = 184; + static constexpr uint32_t AccumRegisterRequirement = 256; + // Get the Algorithm parameters constexpr static int NumComputeMtxs = 3; constexpr static int NumBandsToCompute = DispatchPolicy::NumBandsToCompute; diff --git a/include/cutlass/gemm/collective/sm100_mma_array_warpspecialized_interleaved_complex_emulated.hpp b/include/cutlass/gemm/collective/sm100_mma_array_warpspecialized_interleaved_complex_emulated.hpp index d9275cf0..45b9cb43 100644 --- a/include/cutlass/gemm/collective/sm100_mma_array_warpspecialized_interleaved_complex_emulated.hpp +++ b/include/cutlass/gemm/collective/sm100_mma_array_warpspecialized_interleaved_complex_emulated.hpp @@ -189,6 +189,11 @@ public: static constexpr uint32_t NumTransformationThreads = 128; static constexpr uint32_t NumAccumThreads = 128; + // Register reconfiguration + static constexpr uint32_t GenericRegisterRequirement = 64; + static constexpr uint32_t TransformRegisterRequirement = 184; + static constexpr uint32_t AccumRegisterRequirement = 256; + // Get the Algorithm parameters constexpr static int NumComputeMtxs = 3; constexpr static int ConjSwapMode = 2; diff --git a/include/cutlass/gemm/collective/sm100_mma_array_warpspecialized_interleaved_complex_tf32.hpp b/include/cutlass/gemm/collective/sm100_mma_array_warpspecialized_interleaved_complex_tf32.hpp index 9e0a6c5e..33507c9b 100644 --- a/include/cutlass/gemm/collective/sm100_mma_array_warpspecialized_interleaved_complex_tf32.hpp +++ b/include/cutlass/gemm/collective/sm100_mma_array_warpspecialized_interleaved_complex_tf32.hpp @@ -173,6 +173,11 @@ public: static constexpr uint32_t NumTransformationThreads = 128; static constexpr uint32_t NumAccumThreads = 128; + // Register reconfiguration + static constexpr uint32_t GenericRegisterRequirement = 152; + static constexpr uint32_t TransformRegisterRequirement = 200; + static constexpr uint32_t AccumRegisterRequirement = 152; + // Get the Algorithm parameters constexpr static int NumComputeMtxs = 2; constexpr static int AccumulatorPipelineStageCount = DispatchPolicy::Schedule::AccumulatorPipelineStageCount; diff --git a/include/cutlass/gemm/collective/sm100_mma_warpspecialized_emulated.hpp b/include/cutlass/gemm/collective/sm100_mma_warpspecialized_emulated.hpp index 57bc1cca..1be80601 100644 --- a/include/cutlass/gemm/collective/sm100_mma_warpspecialized_emulated.hpp +++ b/include/cutlass/gemm/collective/sm100_mma_warpspecialized_emulated.hpp @@ -192,6 +192,11 @@ struct CollectiveMma< static constexpr uint32_t NumTransformationThreads = 128; static constexpr uint32_t NumAccumThreads = 128; + // Register reconfiguration + static constexpr uint32_t GenericRegisterRequirement = 64; + static constexpr uint32_t TransformRegisterRequirement = 184; + static constexpr uint32_t AccumRegisterRequirement = 256; + // Get the Algorithm parameters constexpr static int NumComputeMtxs = 3; constexpr static int NumBandsToCompute = DispatchPolicy::NumBandsToCompute; diff --git a/include/cutlass/gemm/collective/sm100_mma_warpspecialized_interleaved_complex_emulated.hpp b/include/cutlass/gemm/collective/sm100_mma_warpspecialized_interleaved_complex_emulated.hpp index 140b7854..a8fe8c4a 100644 --- a/include/cutlass/gemm/collective/sm100_mma_warpspecialized_interleaved_complex_emulated.hpp +++ b/include/cutlass/gemm/collective/sm100_mma_warpspecialized_interleaved_complex_emulated.hpp @@ -185,6 +185,11 @@ public: static constexpr uint32_t NumTransformationThreads = 128; static constexpr uint32_t NumAccumThreads = 128; + // Register reconfiguration + static constexpr uint32_t GenericRegisterRequirement = 64; + static constexpr uint32_t TransformRegisterRequirement = 184; + static constexpr uint32_t AccumRegisterRequirement = 256; + // Get the Algorithm parameters constexpr static int NumComputeMtxs = 3; constexpr static int ConjSwapMode = 2; diff --git a/include/cutlass/gemm/collective/sm100_mma_warpspecialized_interleaved_complex_tf32.hpp b/include/cutlass/gemm/collective/sm100_mma_warpspecialized_interleaved_complex_tf32.hpp index 1dc7b02c..a1f25017 100644 --- a/include/cutlass/gemm/collective/sm100_mma_warpspecialized_interleaved_complex_tf32.hpp +++ b/include/cutlass/gemm/collective/sm100_mma_warpspecialized_interleaved_complex_tf32.hpp @@ -184,6 +184,11 @@ public: static constexpr uint32_t NumTransformationThreads = 128; static constexpr uint32_t NumAccumThreads = 128; + // Register reconfiguration + static constexpr uint32_t GenericRegisterRequirement = 152; + static constexpr uint32_t TransformRegisterRequirement = 200; + static constexpr uint32_t AccumRegisterRequirement = 152; + // Get the Algorithm parameters constexpr static int NumComputeMtxs = 2; constexpr static int AccumulatorPipelineStageCount = DispatchPolicy::Schedule::AccumulatorPipelineStageCount; diff --git a/include/cutlass/gemm/kernel/sm100_gemm_array_tma_warpspecialized_input_transform.hpp b/include/cutlass/gemm/kernel/sm100_gemm_array_tma_warpspecialized_input_transform.hpp index 3171053a..25d5c6e3 100644 --- a/include/cutlass/gemm/kernel/sm100_gemm_array_tma_warpspecialized_input_transform.hpp +++ b/include/cutlass/gemm/kernel/sm100_gemm_array_tma_warpspecialized_input_transform.hpp @@ -148,9 +148,10 @@ public: static constexpr uint32_t NumFixupBarriers = 1; static constexpr uint32_t CLCResponseSize = sizeof(typename TileScheduler::CLCResponse); - // Transfer registers from regular warps to Accum warps - static constexpr uint32_t GenericRegisterRequirement = 152; - static constexpr uint32_t AccumRegisterRequirement = 200; + // Register reconfiguration + static constexpr uint32_t GenericRegisterRequirement = CollectiveMainloop::GenericRegisterRequirement; + static constexpr uint32_t TransformRegisterRequirement = CollectiveMainloop::TransformRegisterRequirement; + static constexpr uint32_t AccumRegisterRequirement = CollectiveMainloop::AccumRegisterRequirement; // Pipeline and pipeline state types using Load2TransformPipeline = typename CollectiveMainloop::Load2TransformPipeline; @@ -412,6 +413,22 @@ public: return dim3(MaxThreadsPerBlock, 1, 1); } + // Register alloc/dealloc behavior might change according to the underlying collective used + template + CUTLASS_DEVICE + static constexpr void + warpgroup_reg_reconfig() { + // Compute default-allocated registers per thread: round_down((512 / NumWG), 8) + constexpr int32_t MaxWarpGroupsPerBlock = ceil_div(MaxThreadsPerBlock, NumThreadsPerWarpGroup); + constexpr int32_t NumRegsPerThread = (512 / MaxWarpGroupsPerBlock) / 8 * 8; + if constexpr (NReg < NumRegsPerThread) { + arch::warpgroup_reg_dealloc(); + } + else if constexpr (NReg > NumRegsPerThread) { + arch::warpgroup_reg_alloc(); + } + } + CUTLASS_DEVICE void operator() (Params const& params, char* smem_buf) { @@ -677,7 +694,7 @@ public: if (is_participant.main_load) { // Register reconfiguration - arch::warpgroup_reg_dealloc(); + warpgroup_reg_reconfig(); // Ensure that the prefetched kernel does not touch // unflushed global memory prior to this instruction @@ -791,7 +808,7 @@ public: else if (is_participant.transformation) { // Register reconfiguration - arch::warpgroup_reg_dealloc(); + warpgroup_reg_reconfig(); // Signal the epilogue warps to proceed once the prologue is complete epilogue_throttle_barrier.arrive(); @@ -833,7 +850,7 @@ public: else if (is_participant.sched) { // Register reconfiguration - arch::warpgroup_reg_dealloc(); + warpgroup_reg_reconfig(); // Signal the epilogue warps to proceed once the prologue is complete epilogue_throttle_barrier.arrive(); @@ -898,7 +915,7 @@ public: else if (is_participant.mma) { // Register reconfiguration - arch::warpgroup_reg_dealloc(); + warpgroup_reg_reconfig(); // Allocate all tmem tmem_allocator.allocate(TmemAllocator::Sm100TmemCapacityColumns, &shared_storage.tmem_base_ptr); @@ -966,7 +983,7 @@ public: else if (is_participant.epi_load) { // Register reconfiguration - arch::warpgroup_reg_dealloc(); + warpgroup_reg_reconfig(); // Ensure that the prefetched kernel does not touch // unflushed global memory prior to this instruction @@ -1051,7 +1068,7 @@ public: else if (is_participant.epilogue) { // Register reconfiguration - arch::warpgroup_reg_alloc(); + warpgroup_reg_reconfig(); // Throttle the epilogue warps to improve prologue performance static constexpr int epilogue_throttle_phase_bit = 0; @@ -1182,7 +1199,7 @@ public: else { // Register reconfiguration - arch::warpgroup_reg_dealloc(); + warpgroup_reg_reconfig(); } } }; diff --git a/include/cutlass/gemm/kernel/sm100_gemm_tma_warpspecialized_input_transform.hpp b/include/cutlass/gemm/kernel/sm100_gemm_tma_warpspecialized_input_transform.hpp index 67853638..c82e084f 100644 --- a/include/cutlass/gemm/kernel/sm100_gemm_tma_warpspecialized_input_transform.hpp +++ b/include/cutlass/gemm/kernel/sm100_gemm_tma_warpspecialized_input_transform.hpp @@ -143,10 +143,10 @@ public: static constexpr bool IsSchedDynamicPersistent = TileScheduler::IsDynamicPersistent; - // Transfer registers from regular warps to Accum warps - static constexpr uint32_t GenericRegisterRequirement = 64; - static constexpr uint32_t TransformRegisterRequirement = 184; - static constexpr uint32_t AccumRegisterRequirement = 256; + // Register reconfiguration + static constexpr uint32_t GenericRegisterRequirement = CollectiveMainloop::GenericRegisterRequirement; + static constexpr uint32_t TransformRegisterRequirement = CollectiveMainloop::TransformRegisterRequirement; + static constexpr uint32_t AccumRegisterRequirement = CollectiveMainloop::AccumRegisterRequirement; // Pipeline and pipeline state types using Load2TransformPipeline = typename CollectiveMainloop::Load2TransformPipeline; @@ -389,6 +389,22 @@ public: return dim3(MaxThreadsPerBlock, 1, 1); } + // Register alloc/dealloc behavior might change according to the underlying collective used + template + CUTLASS_DEVICE + static constexpr void + warpgroup_reg_reconfig() { + // Compute default-allocated registers per thread: round_down((512 / NumWG), 8) + constexpr int32_t MaxWarpGroupsPerBlock = ceil_div(MaxThreadsPerBlock, NumThreadsPerWarpGroup); + constexpr int32_t NumRegsPerThread = (512 / MaxWarpGroupsPerBlock) / 8 * 8; + if constexpr (NReg < NumRegsPerThread) { + arch::warpgroup_reg_dealloc(); + } + else if constexpr (NReg > NumRegsPerThread) { + arch::warpgroup_reg_alloc(); + } + } + CUTLASS_DEVICE void operator() (Params const& params, char* smem_buf) { @@ -638,7 +654,7 @@ public: if (is_participant.main_load) { // Register reconfiguration - arch::warpgroup_reg_dealloc(); + warpgroup_reg_reconfig(); // Ensure that the prefetched kernel does not touch // unflushed global memory prior to this instruction @@ -716,7 +732,7 @@ public: else if (is_participant.sched) { // Register reconfiguration - arch::warpgroup_reg_dealloc(); + warpgroup_reg_reconfig(); // Signal the epilogue warps to proceed once the prologue is complete epilogue_throttle_barrier.arrive(); @@ -770,7 +786,7 @@ public: else if (is_participant.transformation) { // Register reconfiguration - arch::warpgroup_reg_alloc(); + warpgroup_reg_reconfig(); // Signal the epilogue warps to proceed once the prologue is complete epilogue_throttle_barrier.arrive(); @@ -813,7 +829,7 @@ public: else if (is_participant.mma) { // Register reconfiguration - arch::warpgroup_reg_dealloc(); + warpgroup_reg_reconfig(); // Tmem allocation sequence tmem_allocator.allocate(TmemAllocator::Sm100TmemCapacityColumns, &shared_storage.tmem_base_ptr); @@ -880,7 +896,7 @@ public: else if (is_participant.epi_load) { // Register reconfiguration - arch::warpgroup_reg_dealloc(); + warpgroup_reg_reconfig(); // Ensure that the prefetched kernel does not touch // unflushed global memory prior to this instruction @@ -943,7 +959,7 @@ public: else if (is_participant.epilogue) { // Register reconfiguration - arch::warpgroup_reg_alloc(); + warpgroup_reg_reconfig(); // Throttle the epilogue warps to improve prologue performance static constexpr int epilogue_throttle_phase_bit = 0; @@ -1067,7 +1083,7 @@ public: else { // Register reconfiguration - arch::warpgroup_reg_dealloc(); + warpgroup_reg_reconfig(); } } }; diff --git a/include/cutlass/gemm/kernel/sm100_tile_scheduler_group.hpp b/include/cutlass/gemm/kernel/sm100_tile_scheduler_group.hpp index 596160d9..2d8728a9 100755 --- a/include/cutlass/gemm/kernel/sm100_tile_scheduler_group.hpp +++ b/include/cutlass/gemm/kernel/sm100_tile_scheduler_group.hpp @@ -164,9 +164,6 @@ public: // Given device SM count, set grid size s.t. we do not launch more thread blocks than we can run concurrently Arguments args{}; - if constexpr (!std::is_const_v) { - args.max_swizzle_size = 1 << params.params_sm90_.log_swizzle_size_; - } args.raster_order = params.params_sm90_.raster_order_ == RasterOrder::AlongN ? RasterOrderOptions::AlongN : RasterOrderOptions::AlongM; return Params::get_grid_shape( diff --git a/include/cutlass/gemm/kernel/sm90_gemm_array_tma_warpspecialized_cooperative.hpp b/include/cutlass/gemm/kernel/sm90_gemm_array_tma_warpspecialized_cooperative.hpp index 90714ebf..3a5149d6 100644 --- a/include/cutlass/gemm/kernel/sm90_gemm_array_tma_warpspecialized_cooperative.hpp +++ b/include/cutlass/gemm/kernel/sm90_gemm_array_tma_warpspecialized_cooperative.hpp @@ -387,9 +387,6 @@ public: get_grid_shape(Params const& params) { // Given device SM count, set grid size s.t. we do not launch more thread blocks than we can run concurrently TileSchedulerArguments args{}; - if constexpr (!std::is_const_v) { - args.max_swizzle_size = 1 << params.scheduler.log_swizzle_size_; - } args.raster_order = params.scheduler.raster_order_ == TileScheduler::RasterOrder::AlongN ? TileScheduler::RasterOrderOptions::AlongN : TileScheduler::RasterOrderOptions::AlongM; dim3 grid_shape; if constexpr (IsGroupedGemmKernel) { diff --git a/include/cutlass/gemm/kernel/sm90_gemm_array_tma_warpspecialized_pingpong.hpp b/include/cutlass/gemm/kernel/sm90_gemm_array_tma_warpspecialized_pingpong.hpp index b3fd29c0..e1fa1c86 100644 --- a/include/cutlass/gemm/kernel/sm90_gemm_array_tma_warpspecialized_pingpong.hpp +++ b/include/cutlass/gemm/kernel/sm90_gemm_array_tma_warpspecialized_pingpong.hpp @@ -399,9 +399,6 @@ public: get_grid_shape(Params const& params) { // Given device SM count, set grid size s.t. we do not launch more thread blocks than we can run concurrently TileSchedulerArguments args{}; - if constexpr (!std::is_const_v) { - args.max_swizzle_size = 1 << params.scheduler.log_swizzle_size_; - } args.raster_order = params.scheduler.raster_order_ == TileScheduler::RasterOrder::AlongN ? TileScheduler::RasterOrderOptions::AlongN : TileScheduler::RasterOrderOptions::AlongM; dim3 grid_shape; if constexpr (IsGroupedGemmKernel) { diff --git a/include/cutlass/gemm/kernel/sm90_tile_scheduler_group.hpp b/include/cutlass/gemm/kernel/sm90_tile_scheduler_group.hpp index b48b5e8a..d746ca73 100644 --- a/include/cutlass/gemm/kernel/sm90_tile_scheduler_group.hpp +++ b/include/cutlass/gemm/kernel/sm90_tile_scheduler_group.hpp @@ -59,6 +59,7 @@ private: uint64_t start_linear_idx = 0; uint64_t total_tiles = 0; uint64_t problem_blocks_along_raster_order = 0; + int32_t log_swizzle_size = 0; } current_group_info_; public: @@ -244,6 +245,25 @@ public: return true; } + // Calculate the log of the swizzle size based on the problem CTAs and the max swizzle size + CUTLASS_DEVICE + static int32_t + get_log_swizzle_size(int problem_ctas_m, int problem_ctas_n, int max_swizzle_size) { + int min_cta_dim = platform::min(problem_ctas_m, problem_ctas_n); + if (max_swizzle_size >= 8 && min_cta_dim >= 6) { + return 3; + } + else if (max_swizzle_size >= 4 && min_cta_dim >= 3) { + return 2; + } + else if (max_swizzle_size >= 2 && min_cta_dim >= 2) { + return 1; + } + else { + return 0; + } + } + PersistentTileSchedulerSm90Group() = default; // Note: constructing this tile scheduler can touch global memory that was @@ -276,8 +296,9 @@ public: ctas_along_m = scheduler_params.divmod_cta_shape_m_.divide(cute::shape<0>(problem_shape) + scheduler_params.divmod_cta_shape_m_.divisor - 1); ctas_along_n = scheduler_params.divmod_cta_shape_n_.divide(cute::shape<1>(problem_shape) + scheduler_params.divmod_cta_shape_n_.divisor - 1); } - auto problem_blocks_m = round_up(ctas_along_m, (1 << params_.log_swizzle_size_) * params_.cluster_shape_.m()); - auto problem_blocks_n = round_up(ctas_along_n, (1 << params_.log_swizzle_size_) * params_.cluster_shape_.n()); + current_group_info_.log_swizzle_size = get_log_swizzle_size(ctas_along_m, ctas_along_n, params_.max_swizzle_size_); + auto problem_blocks_m = round_up(ctas_along_m, (1 << current_group_info_.log_swizzle_size) * params_.cluster_shape_.m()); + auto problem_blocks_n = round_up(ctas_along_n, (1 << current_group_info_.log_swizzle_size) * params_.cluster_shape_.n()); current_group_info_.total_tiles = problem_blocks_m * problem_blocks_n; current_group_info_.problem_blocks_along_raster_order = params_.raster_order_ == RasterOrder::AlongN ? problem_blocks_n : problem_blocks_m; @@ -302,7 +323,7 @@ public: FastDivmodU64Pow2 const& divmod_cluster_shape_minor, FastDivmodU64 const& divmod_cta_shape_m, FastDivmodU64 const& divmod_cta_shape_n, - int32_t log_swizzle_size, + int32_t max_swizzle_size, RasterOrder raster_order) { uint8_t valid_tile = 1; @@ -310,7 +331,6 @@ public: // Use a warp to "speculatively" check if the work tile maps to the next 32 groups int lane_idx = canonical_lane_idx(); int total_problem_groups = problem_shapes.groups(); - if (linear_idx >= group_info.total_tiles + group_info.start_linear_idx) { group_info.group_idx += lane_idx; for ( ; ; group_info.group_idx += NumThreadsPerWarp) { @@ -329,8 +349,9 @@ public: ctas_along_m = divmod_cta_shape_m.divide(cute::shape<0>(cached_problem_shapes[0]) + divmod_cta_shape_m.divisor - 1); ctas_along_n = divmod_cta_shape_n.divide(cute::shape<1>(cached_problem_shapes[0]) + divmod_cta_shape_n.divisor - 1); } - auto problem_blocks_m = round_up(ctas_along_m, (1 << log_swizzle_size) * cluster_shape.m()); - auto problem_blocks_n = round_up(ctas_along_n, (1 << log_swizzle_size) * cluster_shape.n()); + group_info.log_swizzle_size = get_log_swizzle_size(ctas_along_m, ctas_along_n, max_swizzle_size); + auto problem_blocks_m = round_up(ctas_along_m, (1 << group_info.log_swizzle_size) * cluster_shape.m()); + auto problem_blocks_n = round_up(ctas_along_n, (1 << group_info.log_swizzle_size) * cluster_shape.n()); group_info.problem_blocks_along_raster_order = raster_order == RasterOrder::AlongN ? problem_blocks_n : problem_blocks_m; group_info.total_tiles = problem_blocks_m * problem_blocks_n; } @@ -356,6 +377,7 @@ public: group_info.start_linear_idx = __shfl_sync(0xffffffff, group_info.start_linear_idx, first_succeeding_thread); group_info.total_tiles = __shfl_sync(0xffffffff, group_info.total_tiles, first_succeeding_thread); group_info.problem_blocks_along_raster_order = __shfl_sync(0xffffffff, group_info.problem_blocks_along_raster_order, first_succeeding_thread); + group_info.log_swizzle_size = __shfl_sync(0xffffffff, group_info.log_swizzle_size, first_succeeding_thread); if (group_info.group_idx + lane_idx < total_problem_groups) { cached_problem_shapes[1] = problem_shapes.get_problem_shape(group_info.group_idx + lane_idx); } @@ -390,15 +412,15 @@ public: uint64_t cluster_idx_minor_div_swizzle, extra, offset; - offset = cluster_id & ((1 << log_swizzle_size) - 1); - extra = cluster_id >> log_swizzle_size; + offset = cluster_id & ((1 << group_info.log_swizzle_size) - 1); + extra = cluster_id >> group_info.log_swizzle_size; uint64_t curr_group_cluster_blk_major = divmod_cluster_shape_major.divide(group_info.problem_blocks_along_raster_order); cluster_idx_minor_div_swizzle = extra / curr_group_cluster_blk_major; cluster_idx_major = extra % curr_group_cluster_blk_major; - cluster_idx_minor = cluster_idx_minor_div_swizzle * (1 << log_swizzle_size) + offset; + cluster_idx_minor = cluster_idx_minor_div_swizzle * (1 << group_info.log_swizzle_size) + offset; auto minor_work_idx = static_cast(cluster_idx_minor * divmod_cluster_shape_minor.divisor + cluster_minor_offset); @@ -430,7 +452,7 @@ public: scheduler_params.divmod_cluster_shape_minor_, scheduler_params.divmod_cta_shape_m_, scheduler_params.divmod_cta_shape_n_, - scheduler_params.log_swizzle_size_, + scheduler_params.max_swizzle_size_, scheduler_params.raster_order_); } diff --git a/include/cutlass/gemm/kernel/sm90_tile_scheduler_stream_k.hpp b/include/cutlass/gemm/kernel/sm90_tile_scheduler_stream_k.hpp index 1e206aa6..c874d638 100644 --- a/include/cutlass/gemm/kernel/sm90_tile_scheduler_stream_k.hpp +++ b/include/cutlass/gemm/kernel/sm90_tile_scheduler_stream_k.hpp @@ -246,8 +246,9 @@ public: static bool can_implement(Arguments const& args) { - // Split count > 1 is only valid for heuristic and split-K decomposition modes - return (args.splits == 1 || + // Split count must be positive, and > 1 is only valid for heuristic and split-K decomposition modes + return args.splits >= 1 && + (args.splits == 1 || args.decomposition_mode == DecompositionMode::Heuristic || args.decomposition_mode == DecompositionMode::SplitK); } diff --git a/include/cutlass/gemm/kernel/tile_scheduler_params.h b/include/cutlass/gemm/kernel/tile_scheduler_params.h index bc09688d..09fda77e 100644 --- a/include/cutlass/gemm/kernel/tile_scheduler_params.h +++ b/include/cutlass/gemm/kernel/tile_scheduler_params.h @@ -1635,7 +1635,7 @@ struct PersistentTileSchedulerSm90GroupParams { uint64_t blocks_across_problem_ = 0; bool pre_processed_problem_shapes = true; - int32_t log_swizzle_size_ = 0; + int32_t max_swizzle_size_ = 0; RasterOrder raster_order_ = RasterOrder::AlongN; GroupProblemShape problem_shapes_; @@ -1658,10 +1658,8 @@ struct PersistentTileSchedulerSm90GroupParams { CUTLASS_UNUSED(hw_info); - // Round up to nearest multiple of swizzle_size along each mode - auto log_swizzle_size = get_log_swizzle_size(problem_blocks.x, problem_blocks.y, max_swizzle_size); - auto problem_blocks_m = round_up(problem_blocks.x, (1 << log_swizzle_size) * cluster_shape.m()); - auto problem_blocks_n = round_up(problem_blocks.y, (1 << log_swizzle_size) * cluster_shape.n()); + auto problem_blocks_m = round_up(problem_blocks.x, cluster_shape.m()); + auto problem_blocks_n = round_up(problem_blocks.y, cluster_shape.n()); RasterOrder raster_order = get_rasterization_order( problem_blocks_m, @@ -1678,7 +1676,7 @@ struct PersistentTileSchedulerSm90GroupParams { blocks_across_problem_ = problem_blocks.x * problem_blocks.y * problem_blocks.z; pre_processed_problem_shapes = problem_shapes.is_host_problem_shape_available(); - log_swizzle_size_ = log_swizzle_size; + max_swizzle_size_ = max_swizzle_size; raster_order_ = raster_order; if (raster_order == RasterOrder::AlongN) { @@ -1727,10 +1725,8 @@ struct PersistentTileSchedulerSm90GroupParams { int const sm_count = hw_info.sm_count; int const max_active_clusters = hw_info.max_active_clusters; - // Round up to nearest multiple of swizzle_size along each mode - auto log_swizzle_size = get_log_swizzle_size(problem_blocks.x, problem_blocks.y, max_swizzle_size); - auto problem_blocks_m = round_up(problem_blocks.x, (1 << log_swizzle_size) * cluster_shape.m()); - auto problem_blocks_n = round_up(problem_blocks.y, (1 << log_swizzle_size) * cluster_shape.n()); + auto problem_blocks_m = round_up(problem_blocks.x, cluster_shape.m()); + auto problem_blocks_n = round_up(problem_blocks.y, cluster_shape.n()); int problem_blocks_total = problem_blocks_m * problem_blocks_n * problem_blocks.z; @@ -1803,24 +1799,6 @@ struct PersistentTileSchedulerSm90GroupParams { return launch_grid; } - CUTLASS_HOST_DEVICE - static int32_t - get_log_swizzle_size(int problem_ctas_m, int problem_ctas_n, int max_swizzle_size) { - int min_cta_dim = platform::min(problem_ctas_m, problem_ctas_n); - if (max_swizzle_size >= 8 && min_cta_dim >= 6) { - return 3; - } - else if (max_swizzle_size >= 4 && min_cta_dim >= 3) { - return 2; - } - else if (max_swizzle_size >= 2 && min_cta_dim >= 2) { - return 1; - } - else { - return 0; - } - } - CUTLASS_HOST_DEVICE static RasterOrder get_rasterization_order( @@ -2496,10 +2474,8 @@ struct PersistentTileSchedulerSm100GroupParams { int const sm_count = hw_info.sm_count; int const max_active_clusters = hw_info.max_active_clusters; - // Round up to nearest multiple of swizzle_size along each mode - auto log_swizzle_size = get_log_swizzle_size(problem_blocks.x, problem_blocks.y, max_swizzle_size); - auto problem_blocks_m = round_up(problem_blocks.x, (1 << log_swizzle_size) * cluster_shape.m()); - auto problem_blocks_n = round_up(problem_blocks.y, (1 << log_swizzle_size) * cluster_shape.n()); + auto problem_blocks_m = round_up(problem_blocks.x, cluster_shape.m()); + auto problem_blocks_n = round_up(problem_blocks.y, cluster_shape.n()); int problem_blocks_total = problem_blocks_m * problem_blocks_n * problem_blocks.z; @@ -2581,12 +2557,6 @@ struct PersistentTileSchedulerSm100GroupParams { return launch_grid; } - CUTLASS_HOST_DEVICE - static int32_t - get_log_swizzle_size(int problem_ctas_m, int problem_ctas_n, int max_swizzle_size) { - return UnderlyingSm90Params::get_log_swizzle_size(problem_ctas_m, problem_ctas_n, max_swizzle_size); - } - CUTLASS_HOST_DEVICE static RasterOrder get_rasterization_order( diff --git a/include/cutlass/platform/platform.h b/include/cutlass/platform/platform.h index ee58b054..46e1adf4 100644 --- a/include/cutlass/platform/platform.h +++ b/include/cutlass/platform/platform.h @@ -582,6 +582,7 @@ template struct alignment_of : std::alignment_of {}; #endif + #if CUDA_VERSION >= 11080 /* 16B specializations where 32-bit Win32 host compiler disagrees with device compiler */ template <> @@ -609,12 +610,7 @@ struct alignment_of { enum { value = 16 }; }; - -#if !defined(CUDA_VECTOR_TYPE_ALIGNMENT_16_32_ENABLED) -#define CUDA_VECTOR_TYPE_ALIGNMENT_16_32_ENABLED (__CUDACC_VER_MAJOR__ >= 13) -#endif - -#if (CUDA_VECTOR_TYPE_ALIGNMENT_16_32_ENABLED) +#if CUDA_VERSION >= 13000 template <> struct alignment_of { enum { value = 16 }; @@ -655,7 +651,9 @@ template <> struct alignment_of { enum { value = 32 }; }; + #else + template <> struct alignment_of { enum { value = 16 }; @@ -677,7 +675,7 @@ struct alignment_of { enum { value = 16 }; }; -#endif +#endif // CUDA_VERSION >= 13000 #endif // CUDA_VERSION >= 11080 // Specializations for volatile/const qualified types diff --git a/media/docs/pythonDSL/cute_dsl.rst b/media/docs/pythonDSL/cute_dsl.rst index a05c434e..5ead90f0 100644 --- a/media/docs/pythonDSL/cute_dsl.rst +++ b/media/docs/pythonDSL/cute_dsl.rst @@ -20,3 +20,4 @@ CuTe DSL Deprecation Policy Compile with TVM FFI Ahead-of-Time (AOT) Compilation + Talks and Presentations diff --git a/media/docs/pythonDSL/cute_dsl_general/debugging.rst b/media/docs/pythonDSL/cute_dsl_general/debugging.rst index 0e6319de..759086e3 100644 --- a/media/docs/pythonDSL/cute_dsl_general/debugging.rst +++ b/media/docs/pythonDSL/cute_dsl_general/debugging.rst @@ -112,9 +112,6 @@ For compiled kernels, the generated PTX/CUBIN/IR can be accessed programmaticall - ``__cubin__``: The generated CUBIN data of the compiled kernel. - ``__mlir__``: The generated IR code of the compiled kernel. -These attributes are populated only when the corresponding ``CUTE_DSL_KEEP_*`` environment variable is enabled; -otherwise they return ``None``. - .. code:: python compiled_foo = cute.compile(foo, ...) diff --git a/media/docs/pythonDSL/cute_dsl_general/dsl_ahead_of_time_compilation.rst b/media/docs/pythonDSL/cute_dsl_general/dsl_ahead_of_time_compilation.rst index 67b4354f..4433f65a 100644 --- a/media/docs/pythonDSL/cute_dsl_general/dsl_ahead_of_time_compilation.rst +++ b/media/docs/pythonDSL/cute_dsl_general/dsl_ahead_of_time_compilation.rst @@ -236,4 +236,4 @@ For more information, see the section "Exporting Compiled Module" in :doc:`compi The primary distinction is that, when TVM FFI is enabled, |DSL| generates a dedicated wrapper function on top of the underlying CuTe ABI. This wrapper adheres to the calling conventions defined by TVM FFI. In contrast, the CuTe ABI entry function is specified directly in the generated header file, which affects how arguments must be provided. -For instance, with the TVM FFI wrapper function, users are able to pass in arguments such as ``torch.Tensor`` directly. However, when calling the CuTe ABI entry function, arguments should be provided as ``cute.Tensor`` types. +For instance, with the TVM FFI wrapper function, users are able to pass in arguments such as ``torch.Tensor`` directly. However, when calling the CuTe ABI entry function, arguments should be provided as ``cute.Tensor`` types. \ No newline at end of file diff --git a/media/docs/pythonDSL/cute_dsl_general/dsl_code_generation.rst b/media/docs/pythonDSL/cute_dsl_general/dsl_code_generation.rst index 67180c4d..baeb6a7c 100644 --- a/media/docs/pythonDSL/cute_dsl_general/dsl_code_generation.rst +++ b/media/docs/pythonDSL/cute_dsl_general/dsl_code_generation.rst @@ -7,11 +7,24 @@ End-to-End Code Generation ========================== -1. Techniques for Turning Python into |IR| ------------------------------------------- +1. Hybrid DSL: Python Metaprogramming, Structured GPU Code +---------------------------------------------------------- -1.1 AST rewrite -^^^^^^^^^^^^^^^^ +|DSL| is a **hybrid DSL** that combines two compilation techniques: *AST rewrite* +and *tracing*. This combination gives you the best of both worlds: + +* **Program structure is preserved** — control flow (loops, branches) is + captured via AST rewrite, compiling to proper structured code instead of + flattened traces. +* **Python stays Python** — arithmetic and tensor operations are captured via + tracing, so dynamic shapes, metaprogramming, and Python's rich expression + language work naturally. + +To understand why this matters, let's look at each technique. + + +1.1 AST Rewrite +^^^^^^^^^^^^^^^ The function’s abstract-syntax tree is analysed **before** execution. Python control-flow (``for``/``while``, ``if``/``else``) and built-ins are converted to structured |IR| constructs. Computation inside each region is left untouched at this stage. @@ -47,11 +60,206 @@ trace that is lowered to |IR|. * Data-dependent control-flow freezes to a single execution path. -2. |DSL| Code-Generation Modes +1.3 The Hybrid Solution +^^^^^^^^^^^^^^^^^^^^^^^ + +As shown above, neither technique alone is sufficient—but together they +complement each other perfectly. + +**Why this works: GPU kernels are simple at runtime** + +High-performance GPU kernels are structurally simple at runtime: they avoid +deep call hierarchies, complex branching, and dynamic dispatch. However, +*authoring* such kernels benefits greatly from Python's abstractions—classes, +metaprogramming, and polymorphic patterns improve readability and +maintainability. + +The hybrid approach resolves this tension by evaluating Python abstractions at +compile time while emitting simple, optimized code for runtime execution. + +**How |DSL| divides the work:** + +1. **AST rewrite handles structure** — loops (``for``, ``while``) and branches + (``if``/``else``) are converted to structured |IR| *before* execution. + This solves tracing's control-flow problem. + +2. **Tracing handles arithmetic** — inside each structured region, the tracer + records tensor operations exactly as they execute. No need to model Python's + complex semantics—just run Python and record what happens. This solves AST + rewriting's complexity problem. + +The result: + +* Loops compile to real loops, not unrolled traces. +* All branches are preserved, even if not taken during tracing. +* Dynamic shapes, metaprogramming, and Python idioms work naturally. +* The rewriter only needs to understand control flow, not all of Python. + + +2. |DSL| Compilation Flow: Meta-Stage to Object-Stage +------------------------------------------------------ + +|DSL| bridges Python and GPU hardware through a three-stage pipeline. + +.. _fig-dsl-modes: + +.. figure:: dsl_modes.png + :width: 400 + :align: center + + *Left*: tracing mode records only the path that executed. + *Right*: preprocessor mode emits structured |IR| for every branch and loop + before tracing the arithmetic. + + + The default |DSL| compilation pipeline (mode 2): Python source flows through AST preprocessing + and interpreter-driven tracing to produce |IR|, which is then lowered and + compiled to device code. + +**Stage 1: Pre-Staging (Python AST)** + +Before any code executes, the AST preprocessor rewrites the decorated function. +It inserts *callbacks* around control-flow constructs—loops, branches, and +function boundaries—so that program structure is captured explicitly rather than +lost during execution. + +**Stage 2: Meta-Stage (Python Interpreter)** + +The rewritten function runs in the Python interpreter with proxy tensor +arguments. As execution proceeds: + +* Callbacks fire at control-flow boundaries, emitting structured |IR| (loops, + branches, etc.). +* Tensor operations are traced: each operator invocation records the + corresponding operation. +* Compile-time constants are *partially evaluated*—values known at JIT time + fold directly into the |IR|, enabling aggressive specialization. + +The result is a complete representation of the kernel, with both high-level +structure and low-level arithmetic intact. + +**Stage 3: Object-Stage (Compiler Backend)** + +The internal representation passes through a lowering pipeline: + +1. High-level operations are progressively lowered toward hardware-specific + representations. +2. Optimization passes (tiling, vectorization, memory promotion) reshape the + code for the target architecture. +3. The final code is translated to PTX/SASS (for NVIDIA GPUs) and assembled + into a device binary. + +At runtime, the compiled kernel is loaded and launched on the accelerator. + + +3. Meta-Programming vs Runtime: Two Worlds in One Function +---------------------------------------------------------- + +A key insight for understanding |DSL| is that **your Python code runs twice**, +in two very different contexts: + +1. **Meta-programming time (compilation)** — Python executes to *build* the + kernel. This happens on the host CPU when you call a ``@jit`` function. +2. **Runtime (execution)** — The compiled kernel runs on the GPU with actual + tensor data. + +This distinction determines what you can observe and when. + +``print()`` vs ``cute.printf()``: Meta-Stage vs Object-Stage Output +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +|DSL| provides two ways to print values, each operating at a different stage: + +* **Python's** ``print()`` — executes during the **meta-stage** (compilation). + Use it to inspect what the compiler sees. +* ``cute.printf()`` — compiles into the kernel and executes at **runtime** on + the GPU. Use it to observe actual tensor values during execution. + +The following examples demonstrate how the same ``result`` variable appears +differently depending on when and how you print it. + +**Example 1: Dynamic variables (both** ``a`` **and** ``b`` **are runtime values)** + +.. code-block:: python + + @cute.jit + def add_dynamicexpr(b: cutlass.Float32): + a = cutlass.Float32(2.0) + result = a + b + print("[meta-stage] result =", result) # runs at compile time + cute.printf("[object-stage] result = %f\n", result) # runs on GPU + + add_dynamicexpr(5.0) + +.. code-block:: text + + $> python myprogram.py + [meta-stage] result = + [object-stage] result = 7.000000 + +At meta-stage, ``result`` is a proxy—its value is unknown until the kernel runs. +At runtime, ``cute.printf()`` prints the actual GPU-computed value. + +**Example 2: Compile-time constants (both** ``a`` **and** ``b`` **are Constexpr)** + +.. code-block:: python + + @cute.jit + def add_constexpr(b: cutlass.Constexpr): + a = 2.0 + result = a + b + print("[meta-stage] result =", result) # runs at compile time + cute.printf("[object-stage] result = %f\n", result) # runs on GPU + + add_constexpr(5.0) + +.. code-block:: text + + $> python myprogram.py + [meta-stage] result = 7.0 + [object-stage] result = 7.000000 + +Both values are known at compile time, so Python evaluates ``2.0 + 5.0 = 7.0`` +during tracing. The constant is baked into the compiled kernel. + +**Example 3: Hybrid (** ``a`` **is dynamic,** ``b`` **is Constexpr)** + +.. code-block:: python + + @cute.jit + def add_hybrid(b: cutlass.Constexpr): + a = cutlass.Float32(2.0) + result = a + b + print("[meta-stage] result =", result) # runs at compile time + cute.printf("[object-stage] result = %f\n", result) # runs on GPU + + add_hybrid(5.0) + +.. code-block:: text + + $> python myprogram.py + [meta-stage] result = + [object-stage] result = 7.000000 + +The constant ``b = 5.0`` is folded in, but since ``a`` is dynamic, the result +remains a proxy at meta-stage. The GPU computes the final answer at runtime. + + +Practical Implications +^^^^^^^^^^^^^^^^^^^^^^ + +* **Use** ``print()`` **to debug your meta-program** — inspect shapes, strides, + tile sizes, and compile-time decisions. +* **Constexpr parameters enable specialization** — the compiler can generate + tighter code when values are known at JIT time. +* **Dynamic parameters preserve generality** — a single compiled kernel can + handle varying input sizes without recompilation. + +4. |DSL| Code-Generation Modes ------------------------------ -CuTe’s Python front-end combines the techniques above into **two mutually -exclusive modes**, selectable with the ``preprocessor`` flag of the +CuTe's Python front-end combines the techniques above into **two mutually +exclusive modes** (see :ref:`fig-dsl-modes`), selectable with the ``preprocessor`` flag of the ``@jit`` decorator: 1. Tracing mode ``@jit(preprocess=False)`` – tracing only. @@ -64,23 +272,3 @@ optimisation problems of pure tracing; tracing then fills in the arithmetic. This hybrid “preprocessor” pipeline is unique to |DSL| and was designed specifically to overcome the disadvantages identified above. -.. figure:: dsl_modes.png - :width: 400 - :align: center - - *Left*: tracing mode records only the path that executed. - *Right*: preprocessor mode emits structured |IR| for every branch and loop - before tracing the arithmetic. - - -Why Tracing-Only Is Insufficient for Control-Flow -^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ - -* **Branch loss** – The untaken side of an ``if``/``else`` is never lowered. -* **Loop unrolling** – Loops are flattened to the iteration count observed, - destroying structure needed for parallel mapping and tiling. -* **Data-dependent paths** – Control-flow that depends on tensor values freezes - to a single execution path at trace time. - -The preprocessor mode fixes all of these by lowering control-flow first and delegating -only the arithmetic to the tracer. diff --git a/media/docs/pythonDSL/cute_dsl_general/dsl_compilation.png b/media/docs/pythonDSL/cute_dsl_general/dsl_compilation.png new file mode 100644 index 0000000000000000000000000000000000000000..4689a1b270cebb3ae465cbae7ca350d8f7a43639 GIT binary patch literal 288213 zcmeFZWn9&1^et}8=qO?gA_4*yph$>-NRJ{&qlc6fkuF6-+8D7wq>)r94=LTDNJvT? zIt8Rdy1DD9bAO-zef@uNUtC_yh@$6vo@ej1*Is+=^X#g^70TUoySHrFLMbhEnXqNc zuFqSx{OS7FHvCRk?*+;&Tkda>zI;*np3ZpJqUOObj8i{WpL9|j7kKvO531wuj|Tad zKHb&DH^!FKcbNO8dZ~(q4ez#e>C)tRx zwpF`JTU)zPEVNFQwCO;Bm-_v6uWqy6^1pBJ|NSxK=lJ#YPuO2dTmSoJ%bUIb694x% zTdwZ>r2IeMz2C0K_CMb>JvdAKKi`f16GVQqzrSpc%KiD@ueaP6{QtZh!~c6Z^3nWX zaF}ZpDyph6QBj{ie-?I_xtzBz;!1K0&B24hE-SoTT#AKe_S=75qV8657)*01D=Vj4 zbYvRV1lNQKsl_&$*!T1J`TH}7INCHsNz?ONR@c?FOz!ygmkPI$BdJ%B7#G)Ixw7rx zp+iX;8I|qrx!QSl*RNl1?^@a;ARw@P$DR-A>E>;zgA*ND8F{|~@qY3aa-`|Lh>MH+ z^yyP)XJ=)l%tT?2VscAWTar$J<9KUI2*1^#g9q2F2r6L;fe{g0fvZcO(+Ee|*eV+u zQXOW$6GSE~ouBhsIXdP(dGh4->(`Ny{Ylz+>CC^L?)OW$eJ=OS=Tg$q>0X%VXi3xz z3=Lfy8q>Xg-C?Xb(W-DcM&iYxw{PEuoVUnzoF7dys@qLLvAOm|>~d_Qh~u2``pUeJ z-FU^9FZGk>y(1!uY(^SZ=Et1w-NTU@+xII)5vRKHj4UiH3=IvftcskMtu1!_di~xr zp-5hQaxS+?$HOoss z^sTI>Smc8*OGxh+9|rg!e7q-Q>Pa-}2FqAo%#xy#~`r=0(DKGP=S@|O&57JP*}1_uY9>|@IO z@ZqZ7s1LL3|8ZPQO4qM14}9@`|Ni|{(WdJcALe+q%+)Q1fzfu4`xr3JT`t<_Zc4wIu8I@QU*C z?5P=VPSnKn;W3}>RvMIaoE!czbhig_N5iF49n2?7_qp{i1qFrk=g)uN(cIh|9v&{@ zw4n9j#z&%3Y}AK^gGCYCav51!UQeHXojiYxnct%QmVp5?>BaNsEAb7T$^Ysk2Qr>R zr1<;Enk8%H5H$b#b^lAd$)GsRjZC1R;k2=~PR+`)x3jaeu~C+e)9^@Bid8sv?3l~S z9Ic?uaGXkfLt|q`xtp|dkg1;DF;Z{H*WNH=GXh-vhkx$cwacxSEyhDHE-b8T=|}(I zAsaz$6B83Xy|?3Rr%!X&rnC1fj<=m)9&F>#NPl{E-rn2Y<;ORPo@eXrAJ&g3CS{7x&mF7-Wm`t<4a^ldIKE*>5pHMN%OX@;d`W!>3#P8>Vd z*Oq34ht^%vN@bJd) z;!8urb+(4zIuju`(UoUE`@P2c#OlJtix)3^1@rlYvx9;T=jiU-`2=aGqO5FO_jsRg zc&Lw0LxGb$bxdWNkjqMX=Kf#rb()%7*m|kh*w}OmotM9rz8D`L2k!l`b7x3M$m;TP zSH6Ri=3EuhDoT!=Qk%xOxPh z-Me>ZXJ^Oo)N*nj3pPwDIy$bCodN=#?d^i>?Cd8_h--e@pZ$GsaH2Iuobwos?d{t^ zg0^vS;mR4=*{wE~I!X-(338mr{se@dw;rgN?8?is>}s#CufIp}{rh)hcvfnvPa{tC zi-=6hz`#J3`(JCH9#dG`yRiV_RGOET9>x0xKVI5*N+}BmuH9CFYOpThHx`5G?Zc9yj}F0)fCE+YnyWt z5=t}6^3NSRG&D38rh1AJ667@BxwA+pJ@{7Ow3wuoGc(?nwy$OfJFD~DaDC3*p_fMl ztouGaF85-hp`kIOpjFV(&2N&zpKW zI)qN1yw&roXsx`9Al*S73C`anOHS>I+SWX+HM^5_j~ zE`oH%TAX)IPEJ9A%Rm48V{AP9{#xYsnh+WW2980hzyEH!>m8Hryqp@>ffRI83u5#4 z_Kk{)k_rjEaOaW&Jw5#yuZ_W~K&IUkBr5(5CB~F=_any@g*iATyqJTnC}TO;*m9AK zK;A&SjoW>nRN_?`0CTuS+i)d1)nIMQH58eK@rFRl&TQ}S@Qp8P^NGh?18$D)7}SYk zR@tyHFrYqkC?F&x{xlbLMO|H>S)y){=(%$Qy6fTJQ8@?e&18Nv+)b@tSA}Bof`;q7 zNyF#vZc$-jt-L(R>(|>J?c9eFi)e~%90BN~7G>P&FFPjgJ~!TWcVlfC&=V{kC5O{Y z(JwFaVgmj-F)>Bpu`8JE##_BZLlbswqu9M$AzYZ2nwohKxG!f9TuxL^d3W*AjsnMd zpZJ3`G*cZ}bzNNruInqPC6<@Z#)O{QM2mGop-0CZ(8S6|4>|3CTW^sSVJc- zb(;S;pE+tNU~cE0!)KkBZkw5z0e@&{0`b$$@oI{k?MR}MC;RYTgm5S89!i;71&*Zp z`fF4%ai=s)!+~TQ*R0G~Bt-!sT!4$I`p9ym-q?|}RcV3>aRAW@5;ss*Q&XctpJ_@mZA&!>e*4b4 z9#vSDAx1Hx8%XQs=4P@01ao(Hm%MygieE}vn#glaweGm6s19fP-64gHpd=fQv6kd4 z^R|ZS4>v}0ZO3>cpA4SL(Xz2g?|dWd;(xn}XZP;iKA)DCmt!9PCX#NPWacWWhO0(! zPfSeEpEu9886iwqZZFV@@b=~g00(er;`#)tjtUC4S5?J_KXAUDav7-ino#bvZsJoezbNG~HJ+cw1l^8E@^61g{Gd zvhOW^fUD5_aN~h_V1cNpXkEAnkI3}CS3)k!GlN0g`pOm-=H{VcVf1{a10y4+Cr%`u zTs*#LW_i?Uu)$?@!9S*=y4p{_v7uoHCEdw#a2Lyt%o^T!-vY%5(Fxo}aH(3b0dBek zSBnG+6}Zcy5IWr3djnaHyAy1@Dxq+giHS?a=hJ}$2iUGCAL#w~Xs6rCNGuAeZjr0D zFIVM>0MCZ2i5eL`EAeZDuE%Q`85#aDA)+q1!Ur06S20KnRZs;|e2Mg9Ei`XCw{PD* z+o^`5!0}Bfm9dkaM`L}qZ{N0Upn*-&KTnnKj=jBZp>s~kjA_u6m)OSe)Rfiur5Mki ze22Ruo<$og@rJx{JGI;fLj>ffZLaJ;-b=WhAjDPjkp zFzSA1f{7La=QWQJNaTPpnP$uuBfCiQ(R8s;^J3O74Cko z9v&WETDr%#U-hNS+R`1IN#`401(b)8*Ptr<7=_=L3AJ~0D255n- zxXwcz(=5CcZZtnm7%WKiDoOtexK@?q?LwK~~%us|c_W8e}hNRNsO&2G-G~+mIUqw1Dj5`ce9MI1gJe`lu2wbi~vwDx&uuR$o0f9E! zis->);WsD}W&t|`o@C$Ymz?j}S^HeIAl$2~1py2);4qoCv!=bXK zM!V420qCodVLI=fVm~r6(&MriMn}leD{)6#SB94IN8Yj2GjQMh_qiD0-EOjTWMClb5R<1B#-(|LQ>pzXluYDFPMvb&irlD*6zDuUSiy?O1b7TU(CX<_6nq z&eNy8Em_pmwtf))Yb5&FvuCJ_4VqD5Zp(u-4xHBgscNHk^P?&S4?3I5Jn5suCy zLj-L{$?@y&zK8Z(U!R`g;qK9)A*1%1fdg8AtMirDtnI;T0*B5RDczOH%g&bK4^UnI z`lqCbQ0fZH#q|lq8UqUfQhd^!=A|*Xh)PTK57#N=HWpD&e zBd*XP=4NN(UMen3nQ3UW-`~1jd0Vdiln!5Hvc?x{^bBknKHA`%)>5vn8_1%>goM!0 zP`r@-kznF?voJ8$*B5c8Xtq2u2p)Q-X@)gPxEE)S>UIGq8yL8i z!$6E0UzZAvd8lr(+a%Qc#Fb!9V>VA=rv-EWX+Gnx|4gPpoVoYLB&a>jNYUhYKPROC zslVK-0BEEvwl@8!GkFa}`BixfXdDV?PI_gjRZn5!!|sO<9~x8z5cg-{zccJ|B^sxd z)Yb2eG%%ZppFTRlbl||-$sb^73_^B(o{Wr)NYcO^ckJvKMV*(>*MjF*8LHdaK;0-) za`Dl!Uzr=3r_8mRV4E;Eum1M!8yLRKrhxUp)AmynX(o*&(7`-r8CJD8opLuA4s#ftZ&5cCGPVW@o2X z&K&{&&SWnyFD(sM3IN56XV1JX*{K6Uy3jeE_$^aTPyLimhGY~KnfOfoGBW1T-pH4a zz6Q|BC@4(UM?Rq=#eN@QID4bGCWK#BM&_ozKI%T|;#rf1lY9wRlUu;_GeIN(13f)G zHlH-xqE&m+&sK*Cp=z$Jt>xzn^9eSSNVYY>d}W`&sB!mZ?de2rxdQOsJIOP+`~xQ> z-xsJQX>nV&jM%ibwVghF+WK$bZxKhc(hxENA)U3=uX)7&4l7hly0^Y!%Tw~93iv=* zqU?Y0yk%!=_1TJ4aHrq{>CkcbEjy=YXMGFM^2apRRZ5ea%eCdZw&x4<`bd!5zA#@B z7k69g+v%_$P1n|~TN|PP3Fv^+D$Vc+#J%1fUN2v&X3+5%dYC*lgQ#Tp?ei|`${jm) z$k0+ptgfukMPwHg6yW?(%CvS@rTKC+TyDSXqAn@v-lP=$0Rq=K?FbcaPc~UuSwtv@ zy^)DYTanwQ;FYmTTJXi@_t#FCWbI`TBrAn4eI#FrN6ANtRUPUVAh8!(d&O~~$ZRb2 zd&NB)5Qew_(co2#@Akn$5GXY4EMZE8O+4&p&J?WIDISOTP5yE|_EVwJLRtdQ)3dUe zgzOTYU+Rf9B&tvZ2Gbh80L5uU^=%T)@?sR3SzOG@$Y2n4zNf0%tSK%?&BQd-+k3U{ z!-YZIvb>TswIGxofK--Q%Orj$MsKuEY}1^50h9;C`-`1z7j&~wLEFlEj1qf6WC9Cx zbaZ0UYAXeRl&tc>w*QruC zUggP4^2g2$srB1o=<)b)p9iddD8%3 z91ssR?^QFo({IuWmgVgG_AOPvd~g4?hEENQ zSK;z^`qRxH^u(8A4+h*MbfBAog45h)^Xu}8#SZ*L+UU5rQSf(26`5ICs5+M(TttLQ zN|u4DqvKb(uYB#=o#hz=kSm^&3#BpgA!zC9FO+eMh^%*56`^f&y?y((jLdU3x9`=o z)DI!iRlZKvDbUc;GS43ddeiZlo^%=vwLX?(45AP2^ypCunnRK;Sn(k)9d=D~b90SM z^V3F<7@XYVREQHuyFum2$w_n~{5r?Exl`mr&ZF%DP3BP$U6V$t@)%Y_-oy>KqIrFc z9xFLjgxr9%qE7*%@nRCwv8<^+%bwx1IEkiwifIYl9=gq!2S7~cwWXAHwC#?=Uyp$` zDk|2Y!p8Yn>$3SM*S&+#rI~4d&B9DeSM_l!LfeY+#5pLjdWPz<=+KE%UX2`FT!C?9 zR7~PD_*7D2y8{t+`_7$f^^wp}AulTP$S~l4&9RD1IWj}FVPi)|HCRVj37P_L)6!-U z)9kF=d%ImCLRBZL8X6kf+wmlCr^?r?6L_1@d_-GBKS2qVbbM%YazxDZknV_eT$HlV zv11p*A9g-%%0cxCljZhU+Q(Hul&WMh3* zFf1t{VH{*lKA4xAr`k=(vhzZttw>sGDwwhn(qLjDF{Y2>^l4yC(QCasGSv0+lU+`r zU(hg2N4p~2Ho$+#efhfT8Wm*j5usL8Smd4qoN=Xpf<;x$#HGKKT+mE8+oo5cebANmpteC+Vbqz zuUB)-_!{i-=0FmBrmS^SuZ{8`Cl(v8*xmpPMI}Ppj#PE^_oStzg&q^QJ=y?oezEx# z1z9~nn#5TQzBqVR;Xyph)&Le2>9-nG%E4nOQG~60kEc*x2M+ihCM(%#7hKnthU+7v zEe4<2j5SX{-m{|I6)${%?{re|k5+Yy0=-?WyQo+yb`7%Og$W&E(kst2?9}c#I#xz{ zJbRYoFl(6lU@je%p!4jjS=m3MZqCRJvK1B(Q~Dl4mP000ka>%cYfl)iF}oe%V&4 z#b?TP@AKc%*GBzM)P93v3+46*Mn9;Lpq(ZinTgL=#wR9lAIRB(@o{q_qoLm3uShdc zjlh6_0DpflB$>xyXrzj6_dOhWRfAZ(_0qt?8yD8D<^p zOen)wBqhnc1xGAfN324ahPwKxC#XVXJxp~6avPGNK6-&@88mnzHp65V-r3b<0Ge0O z)TotF|pY6rz-ZQPEn7V?{hOW^+|>H ztL;j`5~h>={gIM$GCa++wJ818Ad1$02gw>M_#v7or0%Jm5dY|=G{Q@fI0MxLaT7|D z(9fSQH&Q{EgqBfpY-T?I;79XxcJ##oqV%5wuTW1pHPWRxkIj2Lf4+O)zQx7G!1Aa| zV|W9^1Pu<4I7ThJ(Wbu?PB;|Kyor8B2hjpWx$a2O8KZ6_N<>6N@DKH*Q$4G5BRrpO zO!t1AsTfcgjLEd@ifbGR7j=s`#272 z-vaBWqwlSscA6U-zxVU2J<0%j7>y(IsZGAwZrPP9T~b*i!{obCacEfpx`JeQq)AQm z<3ZM=F5-?O=&W6>l7Cf#Erq;46}GErm7$orxVV5<@9{0&kd(NJ2t&4kVFC_zQBZV% z#zOaooO<)-H{6?(I#6wYQ-$SerJ*dM`?ko6lX@yo)F{Gw^=!JUlNdD%EQ zyYuXm(L_I2V3iBn&&UXg!ARg)lMy1wbmVK@9Iy!O(`#DQUH)fWZ!xwDqGR*(J1Fo7 zY$BJ!q>RaR=`BjZOdWK;TZV?FUfdP6wE^yg117j8!|EVz=cQ@nQH6X?ULJ-H@dm1- z&qb!(?f7@6+B0PDGM;APXihbdb;xdTG}YJB>%*K2#hjJ$rk?ulrrLiBD;$(fIUg&A zUeNQ{MWPDadsJ}^>TJEoj(`H-U7)5f6-2&X9!+UR8%s;*U{H8Wn9zFh9Y{CX=hqQf zn04>R&C>~GO=QAHsnX5GAH5?e&Qx|E?YUbOP<%WY2pk)6?pH&6Hku)vf!ib`Bp`#+ z(|L1r9ga6-V&)gX0a-(ZA)!#ueE_6^<>bkz*jS7#XTJEdq8A3rczAgDCGmmWAhk$^ zJio5Lm&u(x{@I7Ak%05k2jDP-W6UJMse&)^A3b^lDTC~M@${)W;L%T7O6u(1dZjaF zC^RXxCfeW#fUhIM#d;lxmGY?{(AQ9KG4RZ&5hMx08`S{ z)z#hYTIqM1DNYW7jxz}h3nOW(;}Q<@nlxN~iEc!#)PvUO>ea7M#i4j9t@4r{M4N~5 z7`=qF%Nuz|A1X3RB53>*wAwb~t=z?L#GmY)Q54ujJPgp9?9i&DDso-VO)rB&)}5Rw ze8=G-hR;xeH`cs^tr<=Fssg#Vv~z6`anLqKEsvi(IRFynALIM-?d~cSDA`f&I0c@> z%p(^7kR$pH;>y(t6|;``O$>8^gi(KDU=>BrYt^l#aAuu{pI@!(=@CpM{v3G)t|fZs z+rM=8*mT>otdeKQO}Xp|Sv4eoj7=$-{Ty0;U2)fYs!YP-KXvugPjyu)wvm~I1-kHE z@FZUr`R5)Up?qe|I0lq0N1QxHK3o~oKS+e z`Yz@{O=SK7DoN zMy~c=WJk1cBm@a2o(l+L6tC!yVt%w$%VXoatSA?QUcWxg$E_{0$n(Vn@H2Qo)*z|* z)f3zPGRl#|Wfe-^tjGECF((psQJDawR@z4AjB3@Zvat*yth@x%WGT6FdK)+~30Gn;D*9aiYuP`fw^)I3d9`VG1|J0WuflU~G`;Jvuhmm`8;lnO9d=cj5?e zNJiB`$8;gjMZSK$FiZR+rS{F_qR;H-^71TG62z%Ea%qYXYY8}eOid26%ID7syA}aD z(Q8lQt{Dj(Cbq7M3k++DZf)xSG-S z<#RVbKLvf9=q}I zf)P81_-enGn5%VA0^O$&Z|`F}T~A4ty7}8g`S|#NBlhIkbJg~uGXlx<9r{G|E0&9k zc98oSw^;lT;w=0YrsK~jxVThYT!&8KR+{>#DQ^9%1g3QlF z8ld8s*`RU*AJCLQ=={(6v8C7Q5?NL4VP|DE0yPAY1Bu2s0h(X?-d~m}X`MtMJ>by6(t{xj7pPYB32*-HziK z()6Ywq&L~w>fsL;J{~q^e5aG#B5M8X5ql$kmS3)iI?2fii8x6NVq5=lInB;q&F>u; z7>FumaGZx{7N?(Vrt~F(XjB(Y++QoTv5b~>pdHbU;I=Fpg611&mWRohmR4uZV(2f0 zotvBcO3&Sa3h#9$vU`HVwhhi4l!car>uD<`t{clk5w0sHBSz}!#@8-i{+u9o3mtBM z;hw&}KD5|ho2NryY=O7i@QPUkkZrNW*SgN{TMuL%v6m5NtPu^_%qVYlreYi zU$oJ|3$Ks}L3SrcG|DXP&{Tp)lc=NbiaO3!w3>K|EF)kW3RJ`b!>!hiVaor$=c$a_ z;E)n*Y_6)XPqg24bw>>Tv!Fsm^sWQ$=tgUC(+buCEi7dl7np|vCByl8tM@|v2&k58bbwY%Mi)&U>Sx?;+b@O+_ayZieWnesTh ziShAt#nSHXj~_42?$_4RLbG*i-t3iE9i^(I^cKJ$<~Z2XsYm-S=x6 zVAp5FG?$OXH!a}wXLtg23s#ng1_x>HUgtZnro$4GOfx@oE&0U!%Z9>FzfXVPWgy0q z+3@_zqLLi7&U4JFE>XoCi)MK%Cok`!Ax;^n?BnapBz*67RvFGNiIfdGgD%liU=Sy* z5FzRkSB7-m%OvVt(F-C5^4*2DZh;OxyN<4YSFSDQw%1kkOP^DVL3P|lh3a`e<^oJp zVHfUvTV#>+IL!U|vbwrD9^R?>uU)%F&}}$TK{&OLONVbHHWPG@m%9%WiOI?CihnCl zgqhyGefz9uHsB9xLZX~HopD-*<9s!S)}k^_#Y4+iSr)n86p(uQeG@4uO}O>PUN0Dp zbeDA}xo$&?LrH&Lx#)nRqT>9lNq2tES`6y&;Ls4OSQbj?%Id29*8b-qf_9#2H#J5KCz{JH`#qKmRv-(w7~<$_CXnQ)~viT3`FI4svx~wmuF^kJ8|N~P!sH2 zLFl%0hD5zy>(rntZVGc`LuPx;qh zf34jAWtXyDJ!O)`7!ixMWBcYSuy<$zUP~>s=gnK+Tzr&nzv>XRh$mO;_{_8)mD)>s zb)fr1oknIjwz!xnXv6IRk}AtOZI(j zs7zfs(KDb_RQMQ!!|%BHvDYe_$v^bQxO5H1z3HH-Fvd+P(uY4T@fjZKOVU zFe)YnI)n%diy<)Rj9L=9v(R$}G6;IGZj8MDK6tPxMxGuFM3KsrO|6tfQYn4T=5p=& zi#3&wFv=#KezCz9Cdwz63D)f5=ZX#WHo)5 z9AAyijQG!1E1R|CT>1l@UOr^Y)| z3nCS-u=dnauH{yh=h$ZLups|mo`0E>iM1$6!BzW+OP&<>>h}kk;f3duKO$Nte-J>%tx^f-&nnsdw#jLp+Vc+#{zFpIply_EujGN~`MA_!SlFGB=d4+n< zO&SkWr`ET(w>O`C*bEpi7;AcQty_QQ+YBsnn3Cmtlbc?94*B#OAGew4n5-e!v~eQf ztkC>3gwXHTF0#g`Y8 zdhf?KSPVC`^Z{VdMmsdb6Fu&s-8;bonhPsmPAeS^>(oZ0RA;;N(DW{+Z|A3`-fM-yF=}N<}fmv2US2l-ZzhS@7Wpl_-I_zbE9#5Y?KUeth-{7&1IX-3&E^cnU zPyb?se^~bpd{j*5&4a$Ug>b9uJR>*Xg(n9epH8lUCR`Oy>NtC$l|&-pt})%z`2eC1 z?H#%Vc{YjK4>J=I=eolQkdc|azdDhnh-2`WHoa>Z!zfOPKtNX?z^P4l=umT%G!+Ce zNB2^$!58uIW_flKj*F8b(4Jr{gxAm-MlsAh#O{8Vb9Q#dFv!{y4;K?1P1_>w0Gf`6t!V< z&ZM$76VL?Sd{ZEat~(2bw&hA@qXxt5?%t zN|^nFJ~PO!Y+VwcTghsBY>-7rv)6VwKh#@rY%J;D%i6|8!*Y3l#tu@Q-Bd1K^PM5S zycX}<)bDX;azdHKSzb{Mhkp5$78h=G71K zSab}E&@i#Y4*99B)qK3AqT-+b0J52qk`gSr4k*P?UH2Wi@Bu0}-f>y9=3rPSU=A8C z2Sh_mgdn-D1s?zPRyS41Z`D*Vj1xTv_$w8wf;p%Z4h8-GXiM_c_wUVsXpmOmYI%A2 z(7?d{ux(4Tcz?{)xdUo{hV&E4tn%sy%V!@@)a{TA_g2iS7(JQMwN2bdL*yvQMsAe5 zBuLr!q>86>n1K`_UcV`ro5&{V5*VfWgxIaLM{jhXi(`PtL~T|&PeZP~&9ZresL2(c zAGX3$H1ML58KRA^j}Kfn0#Nqu-`|1}8|F9;Jg{*KUbxWO-kxdI0|9>?b2+F1?T>yX z%`^=;X_Pm07a(G47dXNe@d4T&-P*wDDBa(Gr(#qDiWn0ciBW*Q~o4DG;Of*%GlOeE9 zo%@$46r$MKHvWWxgLP^&V`(^7X;H_J#?@p>z(LH&(2!Bc&Kx!e&`odywCI?SjKM&I z%?wmcY+Aq~Tjf{u`$Cv;A=4Rn+j(#~Vrc?DQbP$P9E&B{#`W(|O|j=CaauDhN_C%X zq5^^Rp=q?ug;F|>F2GcZz&dHX@TiYR&uR7gY5-H=@_8m=)+0^@ z$gFcnhiQp?OunNY*Oyi>^6#iukDFah&t1OwHv?79t5{hQg|-Q`5mgNH1+vNO!;OpR z72(FYbmo;r+k$8K3XN7%rKOT(_qV37%#+= zx^PMJvHQ0>zhjZ^Vu@hs>j-wYq72G(8e3jnE;VS6bEAa`5iFZ5aKqy2U<*~x9y|$N z=MO0<%aeIikTbvhcC~D1kmJA~x8R69i(9v%AfQ82NldhW1rL>+f`TGLVrwtgyZDB3 zu`i<;-)`0fDT{m1T^Gr)dLm);saYr3e!APe^-2e)=yeytktc~Qh1VU++$MkI`ZSFm zRtbrUYjcUE-4R9S)6~-}op?f;kKt9EPg5c(;)f<4d_>TGGbgXEXh4n-+X%rW2|@w3 zUVx;KddS*MR+cW#+10f$llM1j{(w=K90GK-+1~nwS?~QuNzoO-^ode(hAfic(0w~T3IV`J)=+hqiD)gT19D< z@!Uz`PD*_OX}r+ZOPWt^Vs#^^Y4mIcj(L>N^+ii|O^sWt;IS2F1wCTsg02tfLTjtK zgDqfsez;yA>Nq7GuW>`v7qHB9iQfRh3T>#k<$Y@E45&96$Taj385xwjcEQIaB_Z)R z^T?L9=$YqdTwCw(829Dva_x`^w=9~jo*u5E@$E6*JvMJOD_d!=<=mk6@mcU$5~7%! zJ;U)^DioG6PDU0E5eLUYQWmf^WWIZ7@ z6*#MlfQ6^{2_-$*#m=$&=aJqQdrpp5LDMepz@~E>dN+ftYv_&*4Es!`s-^u7iv9(xNUM*Ttu8u<6w*D5esd4M9qD$KNnr{n} zGlhS^*!fUdOZ{^!BsJftjq`a^tFpi<+m@Z>QAqt;a&be8B zORS-hOl}z(8mcnu1XY2j@LHrK)aw&hUfzoNb*2A36S)hr>J|oqva-gaa+4=?vTM6w zRY*u+N8gL#GVVGxm7w{J>s=$vo10bjYM`6U%*-l(G6);uLxH@7UIOA(A`09sf&s!rGvI>rhd;W;~MRC|EPXOG_&hzmmYQ7 zBR%S;p-=FXjXocLS8nuPNf7x4P*WB|gv%`F4yV1X@taW89WpVSmGDuU#N2-{(GqkX zZEs&6d5?hOkKZxyFF6J{1;ZxM)0&zTVCEPq=BA_hO%X#oXKZO;SCH)BhPBV4^}{hx zN-;5A%}GUH3K8sWA^|bc&YLP?;o_Qt`eUt2UQn#8Y#8S_b?O2JEATkDLAYybQm)Y< zodajT2Lh@>3KYF_YFd{-Y%FIk+F%*IWE+Gk=qR(XaQ)8PH*=3poIYctQf?Ki@`O0W zN;I|WK{^5imE>BaQD5SX7Ww`qW;9Eyq^+9zG)?ixoh2E43&gB(-ZP(d-c9vv8H!oU}TA`#1nn%}&jj>Cxa(!*nJanY*v!`t}y zcoa~G_UN1E1Gy@KvGi$w0JRRC8-sI>E+D-quaM&5(YOoOQB60x>x!C69n3l~ ztq+-`#&u&`(r+8Qs_@~ojQ?5qH=7)#&LOkI0U8XEJF1fzU$?uAml&VRFGYTOaV-7eO5X?G}%kX{mS>|KZ^RzyrdI1%x^QLy&e3Q0|}2yj`nF8 zhbw=e%hJ%0h3g$?Cm^y;>fx{0Exb)eduoXN2mxp+2+!Sn_qxsvduM)TVm2ZZ9Jnzs z0oT{n?Ex8sm;!Sxic>{j zgI0qv%Hx9v52l_F(STu)2-zNOAF`UiYqj~+23!c4utUS<0{eUe`6T=W*OR4_^SDH; zjs&0v-YU-bWjGYBacAaY*dwJuh965R8*fyG#8Q>x?x>Gi$wn)vmJq};luff`2=W0k zD*6P<5<*wGx$%KQ0hW!F%uBj#4isNr+`+6BPF&NN9;g(6ieII9@h5qB8r+|S1vqdG z$QNrAGvvX!x97yy3E&T3*B3jTvPh-oo5a2jSS>NBcXIi{shL+^-_J$VvysGb5R zej%ZbARB1F)HF0ejS#~KLLFuubdZ`b)arAt6riQo!1TA%2y;JT37q@sEep}{Nwk&kR|I~54=i~%(byhG&+yt>#< zWY;ZR!2E)tUyo6N12+ zYuj_0F;c|FwyBoK)i`i(NXt&-c>?-WtWnTqlKstp=Y%$GLgdA(*jrm84FYajUm|`8 z6*zI?krcJ=KI#~|)gM~69)d-Jl8GDDBHebAKkhq}@7!?K>mkK1h7a8u@V`Y9AX3Mc zyEZ*PqraGYnQ_y?t|#~;DYh9ztyy|Le8#)YMc!AL?TdcL7i?* z))m7(7?%AY;yGgS26HF2IiZ8x#M&Hxk8@*#AC0u9d7Kf&vFm35MKQHI@F; zBW(&5k%gHV7MW`boN6N&!yTNSnF)pr2qn$vdmuU_%)!Sn4TdXSgi=^o7=uMvLiX#Q zXD7)@Ney=wl+8|jPfW(-j~vR z*I89+Ce6@$OeHB0#JE{(%_g{se0p(ozJe^dKdkgd%iL6hR+^_rh2SB{XT&uOW`O+? ze~&Nk_dFvP#Y2uIaS>RMLw5qu0N-y4Sc(LAaq8q_)u;!P1Oz48_pXu>@eKh?lF2@$ z&GkiqAUHS{WKdAZ9X?h+p^-)nVCGS@zMy&LOgzQ}@wfTm75%k;ENnL(B$#sqV9SiX zIq2ctxwf)Nws3*NCxq=?4NXn5%h>4=x*W2_OQpNHdOLZEMIa6j2JPODCqXMG6+euH zaahmKzC5#j!Q>q(6?~8;ST1U?gyD7=8(e5owdT;k$H}RvkYFby8!7L6`wY zzw{l6drtfy|RgpogrNZh9KKJ76q((}G`a%542J;Tl??xwOl=y#a1zK`kCYdsOY_Op2>=evTM zdIDrD%%uqQUK%aVonl<50ztJRnA8A*nw#xNaVphwzdpxhniblAam%|mZ*JgG!TO*s zU@F!)xj{8hO}2qR#et$qjvu@6C8$#fL~=KT?G*S^tQuo#fCAW3qG4C>9Ui`5Sbdmp zuS#+YDlBrp3&Zl$uqAUK(5%C;GDbFP{yg z3qxs!f<&I+l$RT!s{*{ehIQcxCo8S+y$_)SUX(jv?iFyHd!Lt=*QpQlnqhLMEF?{6 z$Z!fm>EU%6>^7+VvHi?fVLr@azn*?G_wN^;)YkKaZ zm90;l;o|ymP*a^0fkYsZxJOqcBb)i?a*~#_DK|f#Rq3`uwtK4;L`M&TQ9^~t_^tCz zJN;||F%QSO9*on1Cr_Rf5z)?DRm2W4oainvFBPsYbbxr#$RMldi#) z@3!emLGc31i0QB=0~Ijo!jY2q*dRr2PG3kwp+#qRy8-Syt2Om^Jo5cRHsj=zfBKrf0H>m2kaIH7#U$- z5re#pwepEcv143ZR!D0k8&>JMP|fh!0xWVtU!iPVxl)cv9R&UzJ9om(Am+Mu0lW;$ zk%$P#g6@0S3b8z#y<{gC#v{hfa-LpZa2pYeirlocLfeCagX!q#*karB9n#uw#zjR9 z08Aj{|c@w=u(*T!N;-MuY?th8Nxts!#*I6^j%)qC~9)Tqtrx!;4g_|a;ZK^b~ zRBqHNZPnchjg5SZw2Jg-D@G*bALEwu6m^_`yTG`yYk(Q6U39XtRS*+!L|*y7j7s;p zcK$p*c!nCl3}Ar9xP2{yN(}rLZk(>FDl-sF^d6MEcbnoPbI=s6EKHoyb#`$<;)cK{ z=;L#o+87^i0yEC0n3MgwT4{6U)-^GQx@+85l?j4>Tq~Zmlu(u%=VUkc>UwvxJcX)s z3x!-%US5=4!eaP!r+sNVukM-qaoa^~?8hTtGsCKlzMNVYE)x<#a1m4j68gZP5=aVJ z0ZB`Xq=+Mp;tS*(v-A4v2ukdpa*Lf|W_~%%&!Rj2Ah$@4y<602izdvpPv=XYGx+y`h1&-s+3 z;lGC?e0UJ^VM$l}`z#XC&BD(r>2gnGs}*##@NGiY{LY6LjmKHfYJ$aM9-b;H>W*~ap zIHBlZ5>v8CLf)yD6v>}|SaW#Jav3j0GJSrc-UpCe$f;gve5a7o&Wj1w)AX6_WYxSm;d;J4{qomk1M8yjCq4ADk`vT zH~qW;9|M)4Ax{r26z>S=IgWTfo02*UA_rW;#G0Ph`0I1(GdS87oG{Ma(aC96#}EAy zI*zQ&%zX_;Xyg{2CIgVoT6 z9L(1=^XxbnB;BIRV({zw85t9lGaj(}^!;_`Z_7bvv3 zHd_b$hTR=@0BM@H#J%uk!%IwYh#lUPzM+P+ojZ48%ZohPXnKDQZ88)eQTcly?P6In zhkM;<*Bw<<(XKByE;z&`UIfBb>ZqIiRGF7!LCJUa_~IuU|3R_^uRw6tHW%9W~- z9+XEee1+6kTbl$MiOn>o{P*e;H4CSBwm_8DN7zq?)++Ms$R5jszQrcFx?{6~#B?W;#AWEacj1aq9I_~xU7{=U% za(CjKwj&x@Y$MFWP6_oi=hw&qr%AWxXB7V!C7r>C?tm+yA2_wRe%4=R9&w7nJT_Rc z&20~Efn$zl!pbuS{SV*0JMJ2!8V^wzAGB0Tp5IS&&>TK|56%cYK4Ni%V45Mw&mU;G z3h{6slPEs;>gx#XoNg;~#Qdj9k@a6W+g&kUz$6r(k8ciu!TTKaC>{`P_Ng~xn@-}ws;aVHR-W|iUTwCbMMv1 z*R+OC1$u?%mK%u<-}hK+DDL%2oh-}^R!+%FJ{(gU@R-O`de)znSczCAV+31?a$>NF zOi=Fo?zg88UCFt9a&^6^X-+~$_c9hrZ`<1gT9y9?wRe#HrC*idj36UAIxT>L%H%37 zIzr?iYExIVECdI7O-;?*7FfXHYu()L@Kz%c`HxV~c=qfI$OYs_zbjjKx!F83KP0f7 z&SF1K?CpM19p6Wc6{qrg{F%Y=@ZD4Dhj%hvZt+;X=6oc(yquM8bxLBg=!}O8@8s8| z_e@Nme}tEHckp&kyuBkUW?49>zv#*DJf43@w7+<})3P8j|09gJe|g)p^%q{k8vr!{3JuKU3OQ4(*NDnc`~YhLSZBso5~xQQGPmB?uY6sg z$9I(GbS%W#v5Y8-1OUC{CvA;tUKKSO#RU%(oDYf36oQrjTw`{KJwWrWLvio#PGlbjJZNq9Fo|~JV+4^&;F8Gi&3(cQ#rAL9mJr)q-)0xOyncu*-0Xe% zYI5to{B4s?mQi==?r!(YV;q}uOvn?D?8xotYVMxQT~Jlq4B_o;`I=pCFmEf`sj5v? zRvDjkW`nxizqZOmYchX4wJKGXQ`SVwM(0xhLn_bGA(IQ7f4SQTPYvQ&d|9~vz4nO7UWPP@Kft$NQx z#r4MjqU%lIq1^lT@y=46R-IEK(kf-iQnW}Mii|a5W-P^oY@?DQOQB6=DMi+@l%Zyr z6ecw?PKi#Gb%r9eDMpBFS%25f>3PoY|NQ=5-|u;y?Qv%A`}27(*LA(GcSFm!7JIl4 zjSLM*E?ENaGYPuQe~!+)$)oeQSO5%^P+L6|6Gg$n@c7^%5GiPoV2H*K9JynC?>*|6 z3u{u{$?kb1e`nVxQGt<#UmkvU7&H}o^i)=8e1&-@TahIGL?At^+HWhdl2oW5%Tk;d zMUhBmDi?g*8%?~++M7qUNIq}VXi*Jc6DcK{EV8KreJ$0H9HQ#eEq0cbYx7Adc5AeA z#e(q{^Y*0_2N@^Z(rEYhKaV2ja_kF-BfI{wV1p0kaf0i2Bknkmu{uAGjPXH2-;MGX z?dNMiG6Ekej?oqM{iZEIk{MzM6?f7AfY5U^>}`aal4WIVyzZBV>jdr_a^^c%@*D zPtQH(OSg|5E_7fkkzU5%k=~kVQ(~OKp*^Z8rCC_qgE<29lXN=$gx3`Ajj{L5fc~eM z>H{Ba37+Wu^a(m=nO5!F=V_VLOyN5H5IM8-I=-p=n-d=_yVr{yUl&dBw-F`r%E=yU zS=@4Rh-}Q}A>kU9UOdlwt|fv>i7fRti6eaLd+F6%E2~)C5e*GiSdO5WzTS9jKsDys z?J!v%DdR!sC442<^_dZl4fI4eHieCX6~iJ7%x4LOMVe**j9N?^swC_cswiJog|ADE~Yw*LgaN9HnXc_ilHl`S0gd%3GcZ{KI@3 zZ&MV6fWiQY0Vq{e4BnCA@Xv8QIWglu4~-r?*82ACVUXml+S5+*+_(*IqQb@=-y6GM zFUxA#8_he1gpx>6Gte+3d&C;W)H8!LNciJ$DDfGZ@@Y|;+-6zsV+Tg(X@NUWRwEDv zlbMUV`U&gO1GE5oPk?lcEr*Ym9ST=q;P8*Q`eUSFTB-vvVxTUL4j_19ET6;a1X~a@ zA{9+zz8St4G&*}?b(=0bGZTfOPT=s>3ZN7Xqx#h$) z?E{sT`cWcQHc5OA8CRNax3-_S6>O3)JGv~I(%h6^r^%vscS*lG*~PEXWbtIF2Rh56 zDQ6_PHyjulSt~Sl&@Lz5F(>;+C`6E+{5{|*ax9TPq}8W^Sps9oW-Kjh!GX5^7?7Ce zCOP>S35h0JzAO%hD6&0rG&wZLYougklA@w+(t^J^p7IU=24P?@!|okWc~DdnrwqCH z`0DE3sFKlez+EcUDE)n7dvzNLb)cxDqX3Z;Lq+Zqe|6uiW-L5f_UYFU-xeh~mX^tw zo~eSEWKx-$C(mvJZ)^Y$DDZ zneET6Y8w2_oLr*XpNhZ3Ap6Jjx{nuAO-WrO=Llj;3tu zShhL2W3K`=nOUbkRH)?f^|*yatb)c+EAWXkEtWYminK@^d1LW{5hLUDAeIvKz*j{% zR?E_i3KyGs9-WG^EJvJzSx-_PZcF&o$e&vz9e4`()+>xsBMK=ioBi)&=qMHdz(av- zZEK4Gzy529Ccp#Ep8XYO2+&=jdIjNt=}L6BoV+}Wf18S~*byL8I6s(C8b*a#)--R2 z7+3}kAgBQ=(cq#Ug+pE3vvh?R1z0>_tOTNhjEoHQG|XT2OH3O|Z%c8Zd6ndS$X-NZ z^W_lB5Mw!MkGr`648%5}^K0k0bRSt3S86o*6N!i-&u*6y5f10%(PK|jQ-s}qgx`5V z9UKxz$0TXJq9LcL!8a_2tD<6Pn)DahEy7j)gLUarns}O@@Zwp%eOq`OUmvq3GE0{H ztJ$1)X0i@=>%fkQ5rU=Sr?q$au3Ntz(2VHSy$^V3M~)tiDd@1GP>5O0ME1jn-*FT` zA&jG!!t$gV$IcgDwVb>Ix zE*m4kl&WLaaVbbMNz_U*ON;3F#xzKxMb0OO(^WN&Ka@-hBQZ=PG+!PL5!fQYMkH>C zT@&dnztE<-D9ozKux(B>-wzIWdtiwLeBpnB!gMi9P+w@s<$dA~Qna+y)zm=2X4>s% z{VuH7yZZ8&dC?8!8T5{kepLX$P|8ipUu>LyM?dGE@8F#w)gIayY=|&o?_hf`Q zLPSF%PvESJ>PTaHzI3}Nc6XQJ8Wx>H60axH^t^f^ruWKnbv0NgUcJ=v-mH~eBt}bm zoD>$Z7?VH1pRCE?v45mLcEUs#pKNo&qM) zK=?w|#(Fch5$OdN9Ktxf)C{1{M(?IU@e^J*S7AL)&_~I+##xu36*03DU(Ctjwg12#aau{$lQ}Fz%SjU|0 zz<#7ejbqyXfdgnlq@H35%FPo=lzU24`)Je?bv#w-Lo;^FJES0@UvrGNZP8@Sa1dqf zc;}Tc8mLOuEVo_p^!5yrh+<=^ph7%+8P9m;rNss{>ta7Wqr482Xf|v+{=B3l9pn%g zI%?NH*5cDnp40tNZU_YGD^{ecJ)y)H#azAWjgk&v?o0Sc0S(3A&%n2DY1iPJaGS|I z8(QrIAhBp-oM2W_@e#hN5TOHLcv(@WKL>L?fP|ZKoWm#Q(FvnoKoSAK_XxJ)#Jw5_ zQ2-<0e6|D517U*MwAT>4!fmq%Dooy6pli$rR}V={b_9o@!FYEQ7JYjk96t4__!q0h zW6cG( zl4%_C(jNvjhOrwGVzx`% zc|8QHbK>a&p2SRs;^2?^_0Qrdu*fgI0iJ@TZp$*O-JoLx$>-|f(FlK3INQ3qx(?q+ za>KMJGNmnN6Wo7PRjbgEm6o~yH9_S`xK|bwjDY9-eid<+@XdB{A-Y&VcfMaK98EC{Wuk=sekhHk^~H5VhUJ}N%vRi7s3kF)dK#)--3c-6v&Ft zz?vCDx-mkLVz%Y+15r|=bt%l@IV3hER>9+|Vm4_k{}O6skOi(;NAK zr$AU)po%>N-TiPMhLDzRJ+gmvsK*;wKN^W0hY-&-#9nj7}ibVd} zYI8%}V1JA_Q?5tW=j>&s+o8z513Yf*#})+eOt@K&ort}6SE|U1hH*+zsOi-2^&HtM)3(y_QE{@GY(l0$l@K~ zLP1)B6Z1hu7BZaVCof(c07Io7+*uL$je zmsbyG%l038smba`^&@Pu;+rwHYp@1;t3zHMl1d~NGd7ko2iV!fVm9v*8{?8L`e=`gTBTZlfkNJEMF0o^~ES(IIv)kwH{RUO73NY^O2vA_Yy zaZ7b|4}u4PTAb7L=W~~CxqPdLoWKHgjZeY6^|s*shYtioUH>R}f}dxXR*h?SNlGIu z87%U}QBKX7sytIu({?}i63gVfHRW(%`g#|GSQ&Joz6 zHmxIR{ANjH9uunf;8?P%&`}iBeb6d*c1cMLSAjMj=3J2}Y_OtTX6%k;JkVp*{ALrS zd3DWvhgoNAYPQgpq@C!#IJwkg@|piNY1CRpQD?-X>qIlT4;bzDvJ>EG@ElQouNJd} zS1REg3+rHnd_epl&~aa1CT|}>$|g@=TVxB369&;SPl`M&zkXZ8pJZpW-W@>kL8PK( zYyXU%v|83b9upu@(4H+=uy$X;s1MNBfQtylybMr=_4`+2pX zM4yUJ3x29KROPSAn>B7boQ}-s^y?GKqKzR(i3Vb1`yEZA6zZ6wastoHidLXx)E)Y| z?K88gDQ@#n(#v8B&iEKfB&I&3aV;nR+ncn{dfH(sJIeE%v)-oWZ~>EtO=oEU9*t5! zR3S}g4a;gcEHIX=Fj}r?f86Q6dwBkJuA|PLJqu+t^daa;tc-?G(xTwe^WDp3h&Mra z6%(IXN`P7ZW|@i#YUY}#}TOgQoeiouxs zwQ3k5T&lW{n0XgO5&;1Tz*~YR-`OhtJuwIUW1fYPbqA_#FSvaqo}1MDFdbF0P0}-y zwEQ<{YTBY=dPZfdPw>q;3tM%8>%D5SB5EWAZ{Pm@Y!ttyqNq7#Gm7STL{aq5(DwwM zp*u&FISFi@o$cb!05(L8)@*HzqDUoWEVuUl7)1NY^BKn>CdInc#737bXz=~nScaat zg39PDccKL9GJcM_WBw?mQau`cD)G2!o;CS3n}RRC)e4`)Z5U;Vag1YIH)9I4hp9ku6U`m%yvzPAkb?hZEy}_vD)#t=+MjHXZdpTf~&rq zU3>TLh1kV?1D@N(OP4Yl{r&vTo;%m8Dt!JNW{IvSI=puQV^92wV~!fg7seBs2xTsP zv7ne~*a>52mlPX|Z64$3EFl*GRTzG$h3BJq(0^oL{q`#cX&^S2Y(2sRbOEq?sHt5Z z*1c8vo^%X|qH*lKsC6lQZTcXTrV}B9m2-}^_uw{;jgA@_`}^;m7a#^ef_V$}Pc87( zeL@*}*hb%1{aJLzRqS%&~$>O&5xBOu4{O6jach$>IaeY6QvI>*?eg4zI`j`EwV%myfW$vQS*&s!PFJP2G=cN?KLct4&3vmS z>sBX2YBJ52k)84KqHTTS8TQ*uI5F{cTVG^+@wp!8_bGXg%FP-!yM(TJP+T$iS&mhh z2CMlTy2KEHnqGvhYZN<^B*Qc(&xG@P&=bQRC$>^+;{>a%{1pG`QW=@{aN}XIZU0<6 zND8J%8YvAEXUD(eunb31atk6H-k-_!Eov)PyhVn9h~_n{_{_}a%X~w@1hWD3-2-ja z#pn}&*P+WwLBE9#46`2zuay0uKEh*OvgPo;;OdiTLBy7Uxq1&1yVms`SCWznG0gxH z2zHS*S4YRjaNJS;;$6ey{AGTv18}%0As;;XJaCuWl_8td4Cq@oCWQ0Bg9q_)J2oy~ z0Ug5vnLig2>lAP};F;_nU-nBqs;%dZ6JWc$U+13RRCAOh?SI@}Oe_*`R?W81@Ni2? zQq3HAtfxA)jlsDx%q-i}hpHL%DK;@tp<93MS4+1g=wAvo1;NfqrVLF7avCkC)tQf*6KC^?T@94NHiE3j*T9}&B ze!A}$yOmqFCfqlUtjxz~><7e5K*cIbN+u|mX3xGETHT|b5D_sqkwhS5{R5qQcxI6H zLZ)OFy)2Bra^uG6`Kd?KXH|vv?e;bz_Lx~1*`rzjra&w+^1U+K+}%@G_ot=+jQ{m5 zLDvZ^V9ELu@9GyynSKslmiw8+Nu=1v+A5Olnu?~T{iM7Xs6UY0*$Npb)r%mG07zg$z%3X>3(S~qogTH#W3qP4>Yw?F?p zd3AvYw*i`fHO8?Bs8ix)ORcQwLnVIkngQk)Zpu-O$mw1Dej$@5xI!R)qAm|qk7|s| zO>Zpce>)S9#oOXnT=e07%k_H`<5|B9eqKH+Ir!tsv8~;z7iS**zQsFu{>3xuT`#w$ zN(asC8t%O1b-wgVV1ALROJ~9M!M~E$?CP3%{rIuo#|2huOA@9JOL=i(S#xasx0Ksg z*2pPZ84A}&*|yAPJ>sfF5=_mqw6Z9{tL7&c+MmClnOk#fM)tUCTI7a}8*3r@0oDhN zkL7>m9sdA>D@?}dsX_)!o3JVa<&os?zpteC}Ym2DP5;2_PWf&*3#2Xk@K7VeaEOLti4@MxL7#S zERI!w4%l_W@v6K5ZESV^&{2{QJpa6IhpaZ zgJ@zA)Cbfod+MQ|heG(1-56e!l5wHZrlCT?T~%rpDQUUDzvs2Z9?v`521nnG2=X4) zRK~D*LF=VrV%@cK*?t4x<;z=&H@M^qNG-yKe3~vB4;&vdN|3CSF*Cfq!f_-gYG3Q# zAcf%#s2kd0JVDqS-@N(1XN?@sJJM=|qG|V9HIwzy+>R?(nY?;DPgUxDGxk-2%FlV6 zeSk&rKhyG^P-}EO{8STWS=m)cs#vca=X_hLl%0@q+Py43Bky;ES(XyBI!{n-^v>BX zdu-Hw_f4mg(@)0N_CGL0$jX-^$vDR|6>LZ5Z#dAIsbUC%$59vmjZ(jR@!DA*Pkkww zI3aaW4q|wC3<5TSn!l{+#}oVqD4D%119L{zEiKDI>5@LD9U?RhUhM8yW8BO$#4(U@luR??1o}rb+{(z# zUO*5qc@p6R`_&4Ab>kuI-GfHZ$>{^Oy8Qp}3_!;A5}WM`;S~aj5L<&Gv50VU>t7<+VNZ!_~ES{GTpn02Cyf9^cXF^WCM z5={rl*LZ>iL4~_JmUlKTJjzOlVNcTaw5mcIz2u1xGN_ICtM?6Y&8=$ zzv^usqVHV38(16sqL6fHwAdqqDmTb1K=GJlNQ#dq*y=a)Rt+=5w5>P=qXzb|Aj85) z>?}Ei7f#dL{N{Al`N#zbI2x`(f^bD_L6Psc!{Ey9^j)Ku}ON z{=d#Yl#W<23Bl~uf(a9;i&#V}i?J5S3sGKUhB;HHAa=UtqIL%CYM_$@(0>rA9iZj- z`scprkAaU70~N5y0Lp{fPCjy)g}ZZq*S0G&<^`2>@ed8l3W?j3NGX=3(qj<{Vq%Sw z?sx0j+OxiHGn%Mfbg#xk>9{V(q;KRL<+7%wk|A*2_Fh@xGgVU3zE=&t5e^L1KK)F()>$X{vUpei%bA0 zZ+yK)ojI03pk4F=1cCtpnAJJLa|yG#MT%=_nL=PoXgXo^Hc|MbW%&&u17QU8l?Vj@ zCU;;dh_+@lHux0UNEqHpX`gri_zCqeR?XPky)>46{HH#5_(iT?C3kZ{7GKV1kSah0 zkV0H}`)Gzvioxz|jks*>$oi%odrGu*9JlJE%TaAt`kd1R&QWh=by;DcpCfM=1Ke7I z8RQiLTAIVvW5}p6#EK<_JMhn+-nM)Ct|H21gU<(Acum|pSVJ}3KiKi<)495?9}&7` zN@3OmB?@BTgh{+D_C^9TK)(uO@(L6#n1*ZUsR0)oa2tW1YNwyKkU@E$9BDN7uDwpiw}thA!Z3G~1GE`)%EZ+{)5M$EdP$ zatF9#gibU;7l6?S?NgXMCC!31KSTRa>*XwJu4qwg{(YL#u*{D8dVq)t-duy7r^k*& zI9FtvEpVa~pQhMFX`WtF7KJZ=&{`M7x&r5N83Fu=)+RG~_ss7+sdV>$|xqBg|jqf2+jNw?+Sv5+NHSW~-8lWRllydRMmr9YbM7 z&L@G9L6Of_edx%z&54xedS5cpl0@Un+Cb+gT^Nb)?5EvFl_%YKp=~kNyk3{7dKt&6cy16FS2mNd7($zl}7Z;-5{iz#Bb>tVt3v`V#nrsHVyF1iU> zrKgLpK_U`sEYrxDHSq29j6*}4^W)>v8P-03l<%P{E;%?9zdju0s=z%*7HA(tHatma z4McD#4{|}c2CG+`46(+u+rzT`X>-E6o|oAx zw2V|56vbf8mpd>%u$aMC6G#8KBPS=OP2y9KUOo7K0fHW zlp$zFv8${cyEbDI8U)8ocxR%A#gHzVrtd#Zs%?_yrxv%p#fTKxZA?z$BzdxF%%f9@=RYG%>#YWweV~J~8&9r}uM{MbIm&YGS70yH8t&`g4qZ z@ysxw_5PB23MAxkA=tHwmqhu)Xl?A(t5u5wq`j@7Ll{HrZ|7>E(sZR+!7k-c*=KZi zuHS!>x_EBcC^~llO><|>y0Kb(`eDRqB1B>jE?Vp)<(7Ww>a~PU4b#jRS}Y3aH+m>7 zsVJXq8}LbgjHHlFikVHW+@>6*F(f{jcnbV=&E}X_ZtuNKE(soChX~2JpXdMn{~hou-cc(w;i8dj(vLXyHgk#|(lX)Z=8n`}t6D5B4yMZkq6Fa5U=Wa; zCe2AZ8=(;!8Xr3z6ak63xw#J9)FUI!@yhXw2_-fpLCDl!7e50dj!i{o*Fn9`5!S=9 z_T)6uPF)HZ1h)|Y7=SJ?5y?fbjn%|w_SBi*epmi-I~Y6A-UDGEZke~QV(0>MXT;WR zQZhl<=FQ(;T?7Mr&`W@>suv=wRA<5@aGuV&Pdf;X8+hYO zjQO`chAF9jeumYRn3$5M-qFdvSdoEW#0Ta9(LBgS2zmtgY~bga+qA-huni2Oer`{k z^s!^@A^hsSBS5fA=+CSY2U=MVd(%WL((?TnX)du9U_YW5pb(Ly+7OaQpdygKnGk;+$831Ht(7 z&5udj*>LyDb1OsC$>Sn$6{6XU7U<3(I6zcAi0&xL zh;+bp2|=05U}kR7%r%r4s7?s<{$o7x{yBGkXHojl{JwHNBr1rtmzHgpiuq?Ma@vnz z{9QCegd7BrdU<3wrLX&gBYYXg`GBxC#C&XMcRw=&>4-Z3-nVHsYo1%@h)`C8|Cq$S zisskOr*ZbN@Q%QH6gDx28?wOHX#;r_ZxMcN6o? zT_b6b%M;s5ak(+F2_p{7e8G}{*pI19=fm#rNuTN7<9%%o3_xkbT&WzmL(yw$*KqI}f;kayik) z$^$Y$&CvArIb_SVcTtJ4A5V5z!;-lBK8A(BW<2QfB#xp1w1fV{pNbPHTX-e#@c!yE z*7`dtW9bH*3QF1CxtV{3DDLC$#L~t@TO0^Ou?P zA+4;vjy~x_Ypeb@SEwVAQ4c_mTUa;-(FFESfJwpFI?^gJpE?#(0TbD@@G44sRwdyb zSaYqtfvcEe0!0S$5OJ;!yi!fnQn5;x*n$Yb^8wsi^aUIuV%S~SG&P7XSwiw%x?%rs zNg6l+tTdJF?f5Kh?{T;$EWL11<`>QkutEny>H6iy$*Zlglq{_gvP6aviG74cu{qNM zwEzOQ^R~gOO>d$D>-NqLFI%9JiN}%_GAj9PHyq3EP-+a^PorF+GWdsA7&=Cb;?znGnSRtM=z5?FKKu^!RIU!MyT%tVrRb)&b(M`x}7X z&V&6O79*Nm+r)WSo>AirUlQE@2bkPVrqz_dGZD)TDhXAb1WJL%zyEeay#?d3QKbgh zB(4j10dsFvwO?Ue?g&{)oRZBGPpk(X!&Y!Ggv0rd=@MgGoy!*SstM-SSrr{WI1m(u zuyhq4(}4Eo4(XW^=JoME?qmDUbb?w4U8a$Rdf;Njc?~r+8w@;7P@EZ9KZ8@>z8wSc zx(}syiW55Y_it!>-*qtS%q9iM(9}j@a6+gt#7#P0bF3ni@1d+6`A)HyZPU6ZHkV0r zw5_aK57HcHN2?l8#%@^~5Ae5y(n+7ZOInUq+|LO-pL4c>#0Xxz)|ecUNMWs@@(tOm zNItG9A4v>Xlw_9xxcoA*uTq#ddMUtnD~Nb?`;vO{EQpZBb^)Tj|Pa?1AOpD+sl{{5<^+n|(GRqsNkxY{7% zvgZ3BVr>Hk-oXZ~(utKb5#wVp3Ak5fyxN&(3ASFLuR;G=gSF+KNC-`6gJKo+^8BS+ zpa6;c4j2#{*fCR!i>vD8RR_6Hb>T<-7vAsTq0p`int~LZ?s@>t0rYk4W!|z$#+5l> zvS7Sy|9d&y6|-%*Y_v1rI6sDS##K`5K~iskAr|IC!A9-Et7Kdp*l{GN=;CI7Vu%7U z0)ldc)z?kXbk25QEQhMsDubr3YQi}hLTTYy8FF%oT1XnhBY$ODV8GRbK>koD)!D8T zt}vJa1WRQ5&#VrYd0ythUoIscSP3Jn=pd5BdRxMuZtCZ4!4%2o7cc7S8h`s3`#yCp z;_~;&WorU1gUJVnG^fCb+JQ}E)tEh?1Yd=VL z?7UTVo=l?+gb2)}78yFD>^R;7<_jlUqQR;|yb>P7qZy*zKyL+L)Fvdp-opyr1( zQ{apYs2)}t#P|+we;!S_wsq}+F}{Y0#@?0OiqnEbJ|9@yBJ`IAaUhEnnU`+M8wxp> zA@%`t1L@G7l~b>d9ZtMCowhmT??Y`*Hu$QV?$~yGV#RPLb|^D&l<*KnLg@c2j+r)X z2(GmlJ1Ys6HxV*0`Wg!P6XZAew=BP`S5({v=^y4{;k@A=gqdyt7VX(6fAA~_a~`XV zvN@!pYX)gIj`z$KK=o%hd1JH#8$jQJo`bAk-EmYM;q5CTG!GSlvil)@*}=WXU#gmHUpqPHrqJEB2i8jqiUogp&VuIz+rbo* zgk~0yDdr7upf(gTApowkh;Ddxm+Kj+xtMgZZY^+Mky~eGCP{qHNYK@_tocgpZHj?w z?ih|JLa`4B_&&0JpD|$q`hQPw(DY6sUBiV1LXPJPHE)gJcS0DQLS=wEXqV zC1hmQ$2iseghSc|lIPDq|AKPcy&KI8*z79gK0IajNdXcv1mrzD9z}BEmGHBJo69Zoo(k@5@Y}!e* zHW%t6>jpsIqYqAqkFQdL|JiSI=A`MsG#H8+yy<*UN<2^4kE0t#QA+GgbdCj>iQw9X zNgGIw+?*^dsu9l!#V8tOe0=?#JK=$lJU=y92l0QFn3(H($J;2z(=6Kf9_y!=pnVcqP_wi2!8=VT-#Uvv)EMt2>sLckzjNWF+Kuv+hTlcTG!(0Ie1tb z0;|AvJ*ifu<42@;$VPpOIMy_GAfZAYsG8ZdHe1@kwh68CL$PMY^eG=hgt2~ zAM&|+-YiSkTqRVXx)HYZvst`4{y8;9n393%IKzVuvSvXVCRtFI#u1on3xN*-Hw%ow z81uY6MDe3UYISDGhsWQ)6x?|6l=imam>cvx#QmK5(fMs^u3Jh9o;cr7bE|6=jx1+Z zyTR=P*6!Fuffns0{BZDaAxy-=VjCK*R7?>LpWbbkvGNDfJtTejmZ0yhhNc2zvR|G2 z-VzA%`}e&7B+ndImo$}=mv^teso%8In}s7HwPXqQ&eYJ+4+LPM0kp)s;+ct1QlLX>Bh6{>|m9NHfYq$%y(AqA%@+5-p{5a&g`}v@qD0c zPEJS1IWI_D4HJj;rRIiY&e+cfU6yx`^{Y~=8+>`nB#zBWsJZL+)HV1%a?lmIp_D21 zm<_0cOY>-XFl;*_noOkVJoI(h;kq6z4Qg!FrPkidX~Pr@m5PDkgZl?P3Vv|^5@|O5 zK$l?8=o)D1IPsevYp-nsP+_zO8;PAktA-t*m>~2pqu)2`n50!>OoWi3%MzAXB@CM_N(kf!t*r`ntYN+4~A2#kZ{y%eE&T+KiqQDUEP#gculL)m!)X zz$^g3D|a8g2K~&TCQaV);tofQNTGmI2c|o`E@q_9FQpOsPpp$@tgVffI|`}-+Qgd5 zynCsssR%`QvFL!A7s-4*Y3SqLc6J6rY|dGyqcia|bRmkU^o2<6_(H_*B&fu#%F35Q zRqkSpN4fRGhgzH)9C1UHt(!K{*#6GL$)y0*U^<7~bE}a2Y5fg%^kJq@s@vOQRdmv0 zvA_p_wN#o8LGk}c^Zus98f5o&$$)(B$l6+GY%XhiGkwnDyn~H(<>kig)R5t-wKltH z4?hfFyQ8T593!N<>_aTd8ChVuPzKbgyDXP@BB0XeT5fa>4D&GomUHP9^4>M)Vmc)#VuLBQ;gAHm7L1@JfbNl^lkg}3)76!%l`g5wYDSc%z+hL+av_>g{jh5D0V9&qvmN_ z4is?cNFViCY8o^b#wMShpPI|;t;)W+8aiEgYVr4)kgPHqVH@0s-VoREW)U8{J+@}v zuUhR^^GiVu1()e zh*81N!IH9iJ+m+hhyd~iq}^~+bH_NB`a=JPvg&jU7U3-|YHV$rUXv!Qa9rbUlelTS zCrr4Hq&zs;fpCL$SV;7r-@iYCUP;pxODZmITAAtBeDB#EW0ft+e_$@~i9oYpq&J76 z28)bhffHUIrpm#$La7FdZu0FofvqF}BG1l7P{x4fI~bjNEHPRcrI0r;W*as9_)5t| z=akipRPI|;cmF)Ls!D=?w{H$h^JAeBT z`x9N^5XLISFt(38=+AOe56I>%Se&G6Pg(|h6d}Pg2CX4srWKz&b253zA;`3H0`bj& zVZNNgKTtG6X~_w6RujcV3tm0Wdb^=^^rn!!l|zH@zqatWFUq-L*U0}@S%G7zhV_J4 zzKUu31oS^>n}GC#L3r`PWY_`e8bu=NZ?H}l%bHX{zCCy_gk+Cq2Fk6G_0^X!Dhh46 zL}@;Bxgy90#Tre~_h`{Q7tR#z84*PZ>}PNSJJt~}?;U6mqqA0UlVxKDO6NL1TpMp& ze`+-DUjrPuQxjM*1&Q3|Ed>NmK5sKb(1K|LaB@AB9XK+9f0HsQs(&jbGV4aY4dh`f@KH zG0jwk7}lKl485K^LEGg`4hq>qR8Rj>tUUhlgnxSLYndDNTkO`#R3fCo+0G793rQ+p zj}6J0Mt{26=~fv^Guu7-i2%jU%<0qrr5-sm6}Rxa{_4(v%24U(C2fP&aHf$Oo4Znz-*q!?n+m?-rD74vU)FFaen*x`!(6|Jfpc4Vn3DhAN3xHd-PfjS|7=ccXSxlg5@Kf0{QYXtQE{bOkNSsqP=(Jpz zR@yt48e!q|bgo9j4oNM8i$^48L|W*CE@7+gwPtP{X?v=zIeVod9?sOyNjb59cMfO) zVCQH-So@&5b1az+!X4$Zm6esE??gQ642bKIi^1(F4E{9g%ChtdPE2tvC?khq-e71U*E1L_eL+vhn68gTI>!5ZqknrMflGS_#xy~W;% ziDwfbErZ`k+M%H4>5db7Kw{GRY2%3@x6$?z_n$32klfi?FFJ2tB$tayX zo;SaJAQP21>ZbFc4x#e;xSk!*o$oPye0!p*#=*TpS)biPS#B`?-&fzgzMm6+**1Ym zdwQEkVmvF`#x%;`b#R2CCrVLc2xqZsIxXm&jy+pfLVnia7q4J&18oDm3yI2GOY6GB zl%4!r4k6gvapme&D6D~z-Us2C_=j#AMz`A0bHE(|rol%+=OO!`e1|7r)5D{?e1-HF zBdWDEbQYbc_rJgu3&IT8A3~m4W2*9W5S<=W3lptz^rthu&;i^lhWjO)Dlz&M9*GL% z17HdSlLpgYdxp@(1&{Sx^-u@BQY)gK2r7kQK2d_1F(gsSLUN`Y%NKtYCP|HEW*noo zuuRBYUqj9Xi~cBqv?hstOi#p(71wnyu?p^|+C;HKEjohmT_d5d3RNqD!>;3twg2QR zzp1xl=)g`(`+%`w)SVyAcbLQd4DSb3Hb(zn!L1j_KW6+9Ld}yjHMVYj2-Q0SkO;wq zqhUG^M-VX?4{6D}Jg%=V#XF!C(h7Jys&~KsB1!DgV zwa3#S@A`!mFJBTW0R#j*B=hl$0XIw*W?T{`L@SsY3M)v$LU7c1p4`nO&h^W~3^Vqk zG^5QVVe?8X?wP^z%&U=QaRm9i%d$*{BiY;z3~Lq^Q9#4rb@*CBd8 zilVQV5Ij?F7@C=(gTSC}GS~&I<{ou#Bbl40Y&w{T4+s}@pnt??1TKdWlg3xC;CnU- zOyQ;}LM(;=AlHBa65jU+UMeV#Jit=`=?T5v&-EZsz(3>3dG-GC^Upg!z%?4&S7>B$ zZS6N0qyQkersj0cZ$H{a^dU8wiZNwijV_eJPdWIR>7G-;I}&(bPY+kODJ&%W#0>_$T(PQXGuuaLC%ALnK>B9!n{bSXlh~+~6S)wP3qEQ{eYYWk?Twy+s z7R{{dQRSjLnoCCl+VIa(t!`qlK`IQVj9}yw65zV4>)@s0sb5=1MeyYP_# z_BHT)GQrytsA)Dmo6N$*4-b>f&}`w=PGHYEz5$*Uv{pC$_CvG>e1hQ7(fS}>J$){5 z5UNG!GV(A8gex4j>*<{eeEsO?T#!HzK!WAIx3FCZARI)V6Y$c`PV4|PUVIvd!0xc_ zm>b18f^FAe6=uOT9x+m5nD@)LyVKn_u`A9Jx1b;|8>hi|*xY(iXFXF;>y+y>`PplkRcfzaAFP;KREJXy5D*xE&Dsu>ViEB6e(`y__9# z6L%k`-@UhP?3!%}Z*%@g*EV;@+ho=Wd`p9wn$tYi$BA6p_PNF@zWrXU>IQY=&=RY`MXY-1VmfS z#NC)f?EXl2dRv95>a%~IuVJTU|9lpY#BP=$)7P*%=}jxSj&?0|)b7g)3U?iC84S9Q z*_%wsesKug9uB|r?iX(fSeOy^(?=i~@cDQA3rogd=X<97(>2UUQR4V z2HOAt@tC(a)}jGny>sj{N)QwxjytNYN8=L`wqgkpP7H=&2^$f3VL`dx)MTgr3VVEy zDW)E>MAfIvaLJ$@KTm0rrMkMe=&WUm%g$LZ^j@h_8X%HNx`y{PcN2Wy#cJdp-^Y(+St5CV+X^O z?c2BGj!<{+KK%8qE{H3(T;CjsXc-%@Vgo~ zCwZpSJIXwEka|X zCv+x!(6F@p9mXZgRf;Yw$JcMg9bL4@8#*3$`F;2DVby1Mpslsl2hYsC5+p3;H9EyU^eD>)r70@hWuK?VF{!ac``+$*8PoPhO-}keE<1uO?yQEnx zUxl^IasOzFX@u<+CK2e9E-S3>l%^na=jL_$m_`J=qz4M=foT^tJ5G*T+gq^8lUdz! zqtCPI8*hgn`n$1iI?^rhBJ>4u zvDlRr2%!)mApQ`kcnb$76rR6JP<4a z0$@TZD(4lhMNc7@>w|y@QyVK#Eihh2sR{HLF|ZFX=KJ68eTH6Q>y|B$XTQf5s+w}r zz0C9H&m-laV88(^N!B-R+Clc1zXhD0mevXUhR>BLt`A)V@(MowvBS^nWAIS!bYyN@ z$@Pet#bVviTu9bH5#XhV%!unU3@YWNWrx8Wm(zq_((VY z3R6nxg|ElQKX9#F)>F8p62vVKd+~#)7x3G@0K>;I!yaY0d0;coBg{QPSc5TB$D^H! z#A-v#;%7*qi~#-_XBsR1R6GA*T4(`3UNG;I(O#6-ZKwsa4aUTu7>pD)cZc8k>A&T#{~z*ym;guM$&U~{@B)zcL%29 zO;b}55>%%m#96?MDPLgb;|Mu`P`M%j+18gA;q3@*VBy@&-Z@IeR)kF?!gQ=v|t{ebUZocI?)_ni;I6KPdFw<&o%K zOsVDZcD8443aRhT%=za7o8ruY?zQISK+xI^5 z^8#cF)bLAvb_?zh#0o^fZX%^)N(2E2s;ob@VcA9vM4C^{u$8Z5GwPWJ7jk^VP5n9U zGA?5U>1>S~54Kd`MuuDczRUig^xI!D$Z>@fofam)vEZC6R)z*nJR&hX`b`m&Hhaeo za2}A@F77QJF$>8PCS=99PE#3X&=QVvAe%FLPiZ??Bn zpzj|kUR(Md7dYI)aBeRGHw+Xeby`rt^!*f>?|&eu@O@;V|LF0Y{(-xo}mwsz@0A?KHuYZp)l0y_iggZ z$}|sqt1qU<$4fS&B1_QpA*}o`wVRm~tML2C^V+w$JCc|RRuU~_oy1gr?GApPCj%AM zec?mY7bvY>-)^|Z6qy(nvJ)<|rLTm|^c0yuGVG|ZEsN-dc9nk4S+ss-zc~9QT5Pj7zWLRlp2^qY!P0f6c9fLH- z%k=s3aR>BKplo};e#OX22^6$U=d-7VNz|Ja+INDsWH2HcSd0kUhk4y8zJ@~OPQ{Yj z4YDYN*#q-g!)l_;UK)K~)`RV$i&bGE!P(Q@rK}FFNDk!Nak8S9FeM3U#{t{``niuE z7cE>^{LyFfBc1dmFp7B>$SlI}0R39J4i^g{$_&)0h<_l1 zaO?ZS*p8u1Uy4OD(-h=xF}g1k@#M|Qergo)0sBRALXtT#Fkl$)!0@E%Xe zIOvhjsl3u`$v+p(4uLNMG^RuU2oO&C5}V5B8RK1o=EW8*Q2&@mF4>GU>Nh%c?~%;7 z4W17Gb@=#0nSeKaRcD?+M?XOb(6iw*^m&d%t1l8 zW$V`OLqm|@tl7NTs(I1G3a{~!TabgP0=7rVSC%qEv=$#ERrx`x8J)@B%v@z2=o3w8 zpHI|o&3mD3Zsn?=E|f&6U& zy|wbE{uCpQ0!_lsN=MA1tZz|*ve7W7 z+qq9~2cd_+*EO0m3SnOIm4DEPKme5|&tbXc7f5J{-MG2AbAJ2nFs2{8y{`kgfX)s} zaz~9R4C)Qu7V22KXolmZ}x$1Rb|?SX2-cRdBG% zVqZzD{8i=@oFo=5HV43n5<(RiF`&WD2L~900UgON9;EszfM9mAh5VBryHP41LGNoXjgt?$wero#y9TwIU+%hae1d?x%hDTx?I1#NVeSdY>R!F z$@R3_eo5Dyob<{p3*QZAMP8=(8cwm(*M;0>nvJKgSjqbi@Agyf)9tZ-tGIL>M158%NA(uT}MNMoQReUAUo)C6@YZ&ro4(u zlnc?7gTWWW%U5mBC~emXcUxcB4#?QfliJQRNv`GnD_Mk@RixWjVL+drfIr!rnz5^m zmMsCzGT`Sy>ck7INFZH-*CTKO$EzLXCNM8v%7#DcPqP`h8ADP55FBuzQ+fMgkZf?! z9(uu$a6_J@(BDk+_tWu=>@WJ+bO_bCUER}(Xz#zkOmGI6`hW=RZ*CTR{b~qw+x9im!Ztzkf8l35Q`@|f=AwsWML6wxJY}J@xE29C;k!WopmzJ{gV>U`P^CGA{7LK8}I&}=mQ*wVzO^Aysg!u1El z0&eO9As=B0#Q7wZtgH%4g~VKCUHaDwblBh&XD=d0G#MAY+5x>A!0-g=o(D`gD(g8l z;n0}24V}DMsg#iDti7y>kUh%Tr8`5A>vpT*Sg8BAL5T@C!QGSw*g4GHs17isry#}v zaT0PkV6m;7I$mtm1A7#^{7ha`Gd?5)%yd`*B?QH3fRSE?WW%xpXB?Ot;O1oo@R-2a z0Xbli%*@ioksH&l!PC-4XR+bshBIy8coiAV71_xKa`p^?&u_9&0^}h3*%5|Ca_+&z zLWRvzXt)-bpaPe29kw@0h&nG7BN>I&foRdc0q{aSl>psZ5eWdC1tl>AZ*S6=LABTr zvSck|Z&nWNA0G=MP+Swemf*j*QD$%{vhW{}Uea~Eb;j6~1k%MjQjoB>;DdUcF&jIm zT(KF1TpimXK#*vOo31O~n*P+F1afnoPS3cbwm4O;E>W~f$C5(B(z0}p>>XwS&N!CT zK4ogB(w>hvmXo6+^2ifp@)lVOH9+dFz@ z`Dfss18RE|whjo;V0{tkp$ATUP>A*2;1)dqhe8p5f(If#khDXI3(#-?G1rzCncHoT zm?x%(dlx&vVG`B8Qh-bY;?i?bQTN?`C7{%Si3!YtgbtjcxU4%n#9$}|4vFCBOx-;X zVivHxuq!*-+QLI`(8dN+S>Ei?DFP@pY=FmgXR#aT1?v+v^mB2r_J>eHx*g6%!cm9n z0?r4-2BQok*eNFfG6TXh3G+OVY`LFoHlv28fkOl*5D0-qVBwxq9!O=j`%*Olbin;>S{SA-z^RfYN&~}2s zZMI5o*&TKj2nZiMpU6N!Lk(Fl5QM!5a^G&}^fNI4@mRncRQKb(vFzJ`S zg9T)6>8N#Q%@2;i9o4|jGMDeo%FKVujDnW@Xc>YVhcXIePm{I-->hkz4aKY72HcA< zeuP*YeWu`ZMm3kKIxX5wUBag&zGL7_GLrMunbAAvTQ1$wIisc!I!Ss2+#)cJxkl#` zdHZhMUVkGcxkf}!HyRq6m{QJ;#dRHq<*e=;C1y+C)EqS!-Tt_Kh&!%!52GO}0jT_@ z7@O^o$3&W-7;8lPzMFbgO%MVB9CA%W5tXFnvFp1DF&=aKJf;(#ftEy_61K+>6gX+L z)jDsTGee*P0^0x)-c1D_GfsMX@q2i-)Ei*3?wzj_iyPzpK&JX5ff5j^w{{1$P`4!; zn|y3CL_uF`$_^ibe;P)CPWRRDq%))XVk2VULxz`l4~jYo5Utoy-~7tBe3WYS7Tf*W8SksmsNO~?QvU9dpO zs90Tlvek;Z@8WqKVL%^9oxuJ?QPbdR00`p`=!*>`CP?gM(S}h9zY*3njjcXPWiV+9 z@w1?yTR&DdSvOf5D(gT=()qq-m_L!1mnLXm7}1o^$z7uKe6R_{Rt^&)^o@U%A)9`d0EL zKMZc1{;~s)n=K!HCA^#JTyxWh-`mBPM8 z`-AYX!9)|>EzoG1&Php0VXJ9_n+}P589hgJEde&P?AN5Y6snYVfip%xE7KL~5gu}$H5PJq4fT_dLPt(m2yhg zEYsvW)Da2K?O8ZEEsi`JsCsPdXsaSRLJ#QT*dB9mD`0C##TwIg;U$D0Jf@&Xg{B9j z=g=t|fujZK9n4#VgoV}1Us~6@-$ETK#nj_3a5BfL?3uwU6cqp9(7-&$A-v19PZCf$ zGBL^?rNau*GyXcXK#3Yd6x?C$IC@^qvq@TFlIHdm+A8d?zg=;>n}wR?hI{)(lEaG^ z-jE`Il4;Dk8Gu{lXMB+65kspI=#*(C#z`jAJYgG~`QDN!8?*D^iJvA-BAek3c#ng@ zY1;{IL!d`Pn1_xdWEH1M<|CIZ1wyDPzKud12X`7YGoX|MLv;|;MxA&#;}*3^kmy6& zcsS_Z588;JhshWiU^$3kjnou-; ze&6j$DE61!5HBk*jnsVB+Dc;B0e@G0`L z=5&hS zLa77w7;F#VV7Gq%o&-G?7_i>5JGk;9Dvga#o|7+POUmqowlP^a0@BK~eAQ{8!PA0~ z>m@PbdNMg2n*HsHwpiGgA(|L=)0IptoBBSua(SLS0R@P#N4r9D6+hM!2(#_tGjUUv zGcshK(0K>R9;qUDF%DBAiepyhb9-(Xkt@(D^owMV%*fx#y1?7mOreUslE{Bw0)exm zv>B0}44j|)?Ue9Dv^OI@$ZAw@qvYG*kwndkvPVy&N{z`C)}Ht4w5uyDy-t-^xT$DM zoqh3HiP{fQt@Pc$w3of4OOx~2=;e9v6mC9AOO@lfp%CM-oSq?A>bZBUBx9^Fg0R8v zSHe-YUC3U}yTZ^a7ZrQ47+t{^Yx)kp6BS}52U?vx4;1jzjR$8C954GNU#30UWw#+( zkWbI*|%#w}Ij*AtGzx?gj|MS~dsjH;3Iiw>VO02c@MW?4NQ>7G+?EpYfC`vIx zP%`pJfoMz&)I7B5!Y&8N?ma=VP+pON73Gbi5M#-0xi^l+nzm(P5!iBO=5($pqff<0 z)JKwv4c6|}RX)x>yH_W+9(DPLzy9}!OQbT5YP8l)chaDz6}pA1zi~k=3Wn`Ly4&r0|LY zQHq!sMJC&9^F{gg&eID@+$(+8E47Vbl-R>V2wL!Sb) z_m{?5L~+KwabjmN$TNbDdb>ze#vNdApf6j*K9K*oqTy}PQoru&Hr%Tf7M(ZO_X`vpPzU8 zffwHz_4|K*CV2i^wTtg`#bM*(LlTM7MTCXHaJ(AszyJN`*DlUav9>WZW&jF@0ZLG` zfJ%bO;y)q$5)l+xYuN+(;V@YTmhW(Hf+9Ui8ML#L8OQ7dZ$P;0DS~tnhDA{M z8_Z-GW>r17T$#Q<(|>)0@c%=?H(Y8Klb^qFL98j32J8?t=wGnPLC633RglL7M5~0# z9r)!-kTF3q3p+9fOw)~nEz~(T0L17G5dFS?{~LIKu(9m`urQLLiIWae8{k`8 zSy2HTCP+3Ev*c?>E)lXTeJC{ON*^IqXHe7tPQeM&^VZgW-rikQDfu_#v}#={UVs0) zzy}x!&^^joy>t0TjW2eBKY$r9zn_?R|C6Bhn>XwLSgeVm>|lY)0F7ke501c3aZPH+q}>=mXK(J>38sJ0DB@^7&rH>E3B<$+- z-i#nni-YhAiU{&_{gAe#K4J$N4}2+53hv94fjt249ENp*H!fHE${7kQzmgdSw`3U@fNcYwxR7m;Pl77UpP$RWnFISPLo!C+!~{iIgY1sq;rF-J)*OIH zA>lLWO%D)EdISf+))r;W27FTs^ve%$|Mcg3D?m7d;AmxMhk5fRkehx0>l19p$H&G{ z_8c(A38%o=_};yH(NKL(UM412j8VqeSNas-Sis)|gAN$_T7w!6g+hGFf+8^FhdeZn zpmJ*EycBbOH-aJ_U$Ko4OGB~jAcupE189=d`WP3G>|l)`WlV*j1(g8Qgm6B#+Akk9 z(eBS$lfafQgI+Fh2EjcS{w2H{)DX#1kFX1>0EF_s7hu%U%inss_9wV;0~X8h_;EK7 zAAbBmSy4gi2Y&)4TA&LAiwQ|9F2?`)B?)jcPfrIT2(UMS00&BAv~&h6MQ*orpMny-1$If+BjQsPnare)I zYy(<~sw~a?jZ^<4{9XCee7y3GJ3le;*Fu6j$K%H#q4I#wfdB3P{=WC0)pP{V7dT&k z_+LNZ;)2{2XiQKj!A=U}{Ga2CdQ>(;x`=GBz|OUU?zitGSX#=!Yvs?zUU(i$c!+ee zhB^Y(o*1txyU<<)3<<^+o`R+mboh+~0PmJPq(16+z92&vv88$W*@?BNfsJY$15WDnR@s;a7Wx#fR0FdMutH5Cd4JUa>sKzP}RvYNDUo<(A-((cVHfzt(cr{ zftT!!8=EC>A%rcKfyfJ_IR&VOxsU?DUjj}$nI|qa5fDY;svo={3P39mw3|R+{2q^; zsf2`o-!G`Cwy<_xU%Y%7vY%~FzR_T@0^MEsPQhA_EQHmawsJNR0Mb3(*Ks!irbA)K zETf*0x8+OfF?VNK)pPvfkm^Jdsks@3izPbagX8`VWdjc137Xj8;sQ}BU%kQ#BxRs$ z16vFZ9$c}U$vwc61508~C18#AqQG!KjAWbPkAv;YBk5*1lXZZsAf2cH)zEMI(z$~U z2(J72S&c~-bRn4X`vKpBTDb)A34#%HV$iA?0i_MP=WvFXAMFS$P^~lpt;~~Yyhfza zA;0(O6K-zk)>Kqh`fdQe1>F;fL=bK%+5f$HGt@nELkMn{xte7XDDxC1d?pn^x@JzP zb@U#eQ}M|M7kc0g#R$09uJ4}VH;@8!oy_B3T>@ATybK(JHeT{9EJ+<#jAc`~f64@uMhKhgIoM>y z;B{WaNx^aFv>6y*Lqp0F3(K^p;w8yDOlF7w{^pk_8Ewck*vZZ&PEI;_R7{@gliqM$ z59AP6R~e2%7c=22lfqVmr3P6RuxLnvgajDd5N&~@78x1YDr#`CtxnoeFDf6UHy;vW z3HUc;7B>N2R`S^Fe)mDcZ!PZfq`Z&uE%8-0yJzkubTra*%Mul zNgq5I|KL&5&qb4|@jB7G1moX-L9ofnwx7{LNhx8zy<~2@x+0^<`quU9W%}(lt&I~C z$bnL^;h`t0YNKkZ%aMprpK?t>jIOOc=4ZB%97vReEfWwC`PrXl)rtf>uekDn{XeV>I?Mh%$5(&4Ri{vM}<-Q`N{IjcQeYLQ#;8JvY z7!$+t24(32#!g(ohS}U01%~KzX1s*gfd4$th#&==dJF~ess}l|ZScOM6IM_{HBrd+ zM$VKUW~AIUP-Hf}ENkGq!qv32Jr66@XbUAjH8;Y%PYK6yT3LhrF`KFB>#jxi{?wN# zO#vpk`_J8$6x3bt218C%K`jMv`FtlMwWyd@kt?S%g8rnHGd@KKglXFy@VYLaU7Jqk zF}`+uU?^I;_x`SnKCILMehpB3>^vJO@>!wU##vq1%&9Q?BQZ>fDeGOGmGOc(P{JrQIUhABRH}~ zREmTgk00HhR8~e=eI(cFsb%G_p%kRcZPBm{NA8yXa<;H7!eg-M1 z@ygFUjhAX=R4ej7=G&W<5F^KHWeUcd(=~qhO9Oe=Zk4+2_J39>DlTyx$i-mQ&yD9* zDN~e;ueSQ}(L&%{T~WL#P~(YlcS*V@q}a_-T~e7C0Km6=aq>g4z0 z@ivWoKVxwlU4fc)1fwNh3sb;=rlwhSO-@aEHiCa`1ere&&nLRSl;W%*4MKc5Q}uTP zNl%UcA>eb`94aV$uQS)+Q#peip_SusKTKQyT4c7jN+%RXkoB1E>*F>a^DK z8JpCN@u48k&iF)`w7=I#!KR8$G0kIA-?-c;mNz4b&&EG|I2j*PCnj!q^9BYez~iAw zt==~(PXq|W5c)vy!$?r@2gq`vh7GP>YJQ*nF)*%xnghlpQH>R7zd=|7R0LvKG1P7)M4E)mm;RkIA_!pYYf6a;e0HZrMlDQ3_WF=l%>06AufWjrgLR ztaal?qmu$IS%Uk#gt@wU1h*^q@hr0K3&E)qDNdN%$-zvD=q~-*Vz+32|G3BX^P!4u zD_cs^x?}OD=Xg$-|*aJ7&_V= z%+CihYKigB@Wz?u!P?rxlkF(zV-P+09YD`uy6hz*hlc)TsBFVKPeasZMX~kXFk&>- z?R3X}mowst(pHhuN;JP=*GYlyuD}x&MnVC%-dvR<4MEK!15pN5)XLiKRBi0=WcMiv zbMngx2?R|RFGVJ}_0)0a2|wrMwR&#Hp0=c6FUuF;-&d!q)zI{lY3+ygf-eqE4*K6t zS3eQ(>CO0@ zTT^#_*Y(IWLEt=+gr`s{og&-k$l?te&<$6#@gityY0p z3SA*^6|HsH{P5xe%K5c{nsCg7ngHmCXTYL^9145|xOimL zo3}cl5gPoevmV_yITI_D&BN?_6swZm+RgK$l}*><(*5nCqEe2Yv3W&wRI^Cn zk-*=ZtLMIWe-o%|^{H@ZbW98pID0``1}ZiXlfl~c{5(I8;N=t8=r>vTaLpc*lp!!E z|4X=t46g9$j6(dEFOJg4VGLi|X_n z4$J8-*Uf7LBX!OktBQgT^=cY)v*jRWE_t76kQfoE{&SR}c~sL2b#+ z%XH~@~OVe7#oeQmiVW&6G_!}2!9)uO*-o-a82Cetj} z(<+FGiAc01&9#oA{R1sJs7nIK1kRxCtBch7^QiCkQN`W8JQ%cxgaFc9$oA>cF%XQ5 z*X@>ZYn*IlTzLJCWO0T2Bj~`icdW7`uCDI=`)82K!5!%gO#a6|TCZQa4_0%*(itW~ z-U`S+T*-9nVaA4wrY1bPEx}r)7R=&eJ?G+l!ZNtK50qZ&=C zbz3Azk>DPnCSZRHok2*OnYS+9(3L9=IB=r7euV@1z5HzwsFU=>_N8gLY{qhB;R?&x z*2&Y-a_+CK3}h?JP)P!D5%hF`oC2qx?wtMNIsQ~_@6Z;Y$_1(l2Pdcf*V1?aW89ah zZ&8y~*w8R!yCUHqfPs#2!4)cnA%ns~mP3n6;Gho^weAd1b^#X^!<%`fcpSA0jmxQxIGSz z$v`r#&*7=(N)Eos6S@(URQe4K)!qRJob&Qm@!>na7rU>E24*h*=$2)Vprf2~G&LPm zL_~v$4=C40M}vYlipaU8b>rd{WnM$AI8CZi7Er-iA3ruCy8YIu8PhyMOvc?L0vwyX zygFer^6=x0Nc;6Zi{D8Kn(AX^S}u>xx)Mc2gO&xv`;+Lz#cn3>9fyeoL>noqsli-Y z=hh@(wy2-W&xgGUz~xU0>?Gp;o%8QOZEc zdjd|5pB6f5mZT=i&$23t8aFqM;C|t@N?jy#kdv>cQLwTC#K62gn(*dL;e@?4aU2|l z=#IVLEq7!nZ`?RiJ9SlviEw7s(HS{9=Hz(vWhh^3Q^47w&pH33$RGognC2+Hs>&J8 zF3j5j95h!)aPfq{!T0yC#Andh1B?d1>c<=`%cR67u!NlLv=K<;Ao;_iG7AeM?QYsf zZ#!B-s_rsY7P~r+bf+8|bW~4Xkp*K)qPzQxi`^!7?!fv;@V$Ln{V`Xv!qEE2Mqfh0 z(xv*0bRyHYBTjs{bPe(<0}L{^*TaRQ$NL5eyegr=zo$CMQ~77H`B*{ci=LGgMj)b)aiH}g5mMz*;GB=`Blr+FK0!S5qC&P zW@ma^zTM9i^5U)W3(QudJR03p$nnR$i3)4<^ia}5v4eEKOsRh5qoxu(Ua70QKzjo! z5fkfeW-oMFUVEu?=8b&jH&$EiHO|+IEzV!2+=LsCV}4#kwFKLF^t7FG@ZrOo9r68> z%aS>~v`es8ehwf6U~8HE`gR7CfPZfsR!b*@5pjgCFmTh=*0~?NJ&5JXp^~gDG>JTO zGuFxu52p?ZVfDl`>q+%V66_wEaFY)ro$UN#xmvY>U1nMVn==9D?c^ip^vpjGoA`kz0i>+65M2UQQC zzFS&Y7aPNY5p^nWqZn1~oRc)THzW1i-C|i6*3!SvH8zr{6eNqk3I^}WtO84jef)>B zVNC(|xJHK1m&Oj8E36@3WTofWW3ER8US3LqOotlNsWZl~v^-L}anyqph7>L}@l`wg2loq)8*WKh%iKfpTU~8ZqLTb+`PxGOshIwP+bYMwWMMnqD^@@HQ0l~fHg=OR6 z;w%7TAF-Xx_DTof`jVQBxCmI!=kwUj_H*B)laz1)N{{*Wd`WNmCqS&mCeG@}+)-au zmzJffeIvk|eq8N!ig zbwT{td==88rBT#UFM|9!O3n662ld;ZZ^#bpTyl692-$MmCPa&od^=n0&Ac);EBdCg^YcqeO7NTYy^sgsjoDz~YgZS0Z@LDs;KWDEs7*)0@8G0VoftXk&4z#dDt)$J zXLoisxWCprc}^5T#|0{pJ6k?DG0{fY2@Nz-T8_2*Nlusj#7MlfsUw@8vN%5t=@rV{Jc;mlL1S zPEEnhuxODZCe`5^zz0#uS~PH#;C5H}%P?sob{1J0fPntypB-CsD|?hPbGJdvvT zDvD5fx(j{lh0lzV{{Bt7>v1@DGtRLsRE43@p$WC`4?&o}<>A?YiE02SL=c*%Yvq>p zmG*Z<=kzQr%=M`-`}<`e0A)-E4T1`Lqpptx^zdn%NU$kgAZg8tJ zZnWaH1Qo7sCSCx=_96q_%7oc4t+Cu``^)7Hlb7uVrFBImi~2jx%8 zgz(gG(dc~J;f5#~;NrU$gLyiii{>@_{Al=CSiq*YgjL_W1AlQTtBc`$Wn%`X!fs8i zn!VuV_u(RD9bNa~TE3y?=5OFsYt;Kaj=P=Yv5sJ!n}?84J-5wmt>&F&Tf<=%BaQXg z@^Yv=#k7~(`J8w5U-l!gPwB|1sQM6-|0| zP*Fu~KD_1jqc1h}fKiQjbuBTh4g8U_b%x=#$JpJ(>OY{vs z9b9L8va*BRIfE=uy58#EK{5|9oMRv9y|M&%{vxB^_dt9At>g;_ zsGkFBoN>y`15}Ijn@0IW}HS#0b zGK`9!iAhu`m!?Z4qnP5s1G`fvdhymKPfs*YPyO_C*U=I`uycTX9&-3C(=21!urTJe zp{bQ+V?g^#v%l;e@&D*E=Bxb83S=x3^D4twHC} zQm>hr={+KXAyqXCkh!MeKZB);=h12KT6LR}&JY|+MyEV?@yX$G$eGG|=d7H;xSd$ZV?oYii$=r9k z&rcTw$XvI4U=l&F{zcO5V;g%L)493+3?Z-0(_?19J-4T(dPuE0W23nATGYIJ~3Ce+plKZ`ldR0{aI(KyXOEZFTN+O?zP?w=nnI^PJB2A?*Nmjzsm{{P8= zHlXf+FgBl_Ue9gs2D#eB4sK|WBjSXA&dL7xu+aDmT{bV*Arwzq0R-J`1B%*)LnUuo z%Q1+Gm%O22y>_s9ZKkQBs+@p`kR4u6OV{r|GDwnf0c79KKTKw+zF(CTe0O^67THc`1r4! zIQ+VVZNA#t9jc1`3h7YmgQKw(y*DKw1u@;wK!5F;w6TeEwtNCR8?m&+?6w0Hi%QRr zPwxFcJ60MS!%jhLWQjH)U?1*#5aPaTZOuG8AK&paTSB~bw3^HHaLaaTid7__P==06 zoyn`-vb2B}>da?v<~qI*aQMyXys@PfNZ^&;N){GUSI1?EQfDY-1hApnvNZqMB)WU^IDizhez7e@W=bNg(AZ^J7Vl@T0d`# z1{&`ANfX`kEj86c&(MGGoIIo!t+8~r-mcl&8j0hP zgX*8pdcLr$Ylx7w^tlM>@Yku;TG!dPm23aowgV=tNfio9Z$V?#zJIV5N;e&@C z(8JeN61nSAL|Jupn~x8I8Id8FU%%cLk}0+|86;QeUU5abR?76XQt>{jL+TNikEuuwB6$|LQ)$^71_%1F5s7 z06n%m5}rM4h5jO@6;HB(zTodmSrV>7-(PCR#SVKy#AnvZBb`!U)i?F`e)84MOkAXC zypUzQIdn-9bLx1exu4z@iKfR|s2j7;K8lYg%>UFW+P8)0lgqtG^T{)rI=6=*hS2T> z0$cZ)=S8Uu9!#8q!29ev1`&VM97B`rQgU!GHTTNNnX)Du$7c%al_F)G5E|p2)Y5g| zv<>N4Z;V^4xIQ0U&L{-xHhg{A@MsV7!CechsWD)Gvp}Bab}H?YwRj5*t+vNS~(gT0&YFNx*EOY^wPvb;C3Oa7772Vu}Q050mZ|R$dnmF=6s_R z5lJF0tAlUvr4$vheP(rCPis$)Tyry;=QdNZef;9a2M)Ir`uh1YbH_+Y#yh`#`!RH& zEche6<%LjE@OK?T-^QCzR%V6sj)E^*U~h&N8@onZs~Vrl=Pdc`jC^U1Mvtt4;^9IAkTYucSNWnG*PoI=ll8VhnO}>vg zA|ii{moM}#CYh94rOlTddcVeHUgIqRGIxgCH&3r)h=G8E->Rvkuw>xAljUaSGmn

JsKHl#H|Rq3Dy!4ZO-N4s{qcSK{`g|OQkK(K^(x6VZq63&?ob(tZ= z=SvHcl8zj_xOuK;Q;%lK)Q^v~9C$ZOyE!@OT}Lia%`k+$7+YNn!)K;S5%ZkxJw5)^ zmt>8QknAMk9omgTo*0c3c{S^VsE9RLpB-#kHwUmP(TtH0Q~482>FL%_zxAKn5?UFw z*OX8*nB-n$v7Osl0$w5Z`*-Nk`U&oNnFlP?8AcnLEK%;gDla6d+9E!D$Q@Fa78S*h zK}0FdGQ03k4))uM>&-@;9#=61nOm4sVX1)DSvDbS-t2-LpQQ?6$?9lVaGMB~M}hHt zTf%v=`{Cw$FMmcZli6_Z=O|q%KKC>AW1|1ab5|z*eM5KU>Bc^EsE!X0X(c5x(_-Db zb@CSi0)%hea6Q_?^L;)w(`1DAez};PQ>(2#j*y^5uK@IIkbd~KHSq}F_18>yaaUCf zbdnPi0%SqdWtR{g-%A(3C)?Q;4@*n=^N}}RNSt`&PpguGJHHRZMEwth4TdMlZVewo z1b4I%ZJgO0GUXW7MvIZO6G^D$&u;7LQfDe-XSFUXBLl@BJnYaqqp+i!^;kJM_W8=) zK&W8~V~l;$9b3yf_tg2=ai-ZKmeT@choH9$q%jc$p-sn?5SN{+v&MJu8DvCs=~n1qK6r5Ddp$$zP#z*`^sFpp5C3*)E)L+iIlzRX5b@$af_$IwQ4}khj&^3CGsFi2+8&E zSfuID+|h9mqtCIsPd+TZBk7irRiBnGFNU)vXsN9&AN6w~r&X23*8K?(Es&v)Cnh|INrbL~t(aUD6RYb)R~Jm)j#89`mZQqBbP`IIWgzCX&thAJ4?LWW)s$KGM=qlDs!pQY?X9wlz zC_J!U9qXvFvQQ;g8xXibd;Z+4;VPYLUhPy>dKN|%;GG59$n4tM5I8e``ZhxDh?S?2 zqe-4gs#l;a(bV41;*B1MNSg~x6W{n8S=L=P_*UE+$mrUACu;aCyh-XF->#&g4-N<7b z3<=xMkPV-lJh6FC6@i>yXyZzVIaW59N};NgSF#B3#s2r7KC|baeuTuwD=aO|4v*O8 zvF1uSi?>FkhRA8X*-Fm)E^d*Zyur!Kc>moh zUW|Ep90H-#5vx*dYa(oAweE?L^Wz6EG07Xfk1DxI9UX5*OS3X=4o}G$8U0=GrC!I( z&o_ItX0)f_<0>!L-T9l>6f&(Uy2Gf9#GGb3Ke~)O39&Cr(gvY8v9zB5oBp}NNHtSl z-oQ=a-gv&sI)igv$B7~3*x;^wEddTM*1pC$ZOMRjj^!~p{|t{_No#2tnZ^Q~(j3KF zIcn2|Xv^>KM?^&tJ$(2#FK_1Fr}{hFhpZxjC#&E7+C@{kbxX*_H444yHg(1=(ZCah zoB_AITnT^Q9^FSDA9e@7zY5>8!g=l|j*m|Q3GWMqcnpj+Mn*mj!d)Yw`T~U7FdEwO zCr>C?pk2M}{M1-bBX0NCGyCh;(ZXn1+0;4h*WSJkcu@ zWRo5LB`ijsogJZ}_4!UvipO(uhLJc7_u0@gTP%gsK79YKeyV(Q*i|g3O}z32IkTjQ zzxbr2otZ=VfU5$MiQPbdoNKfmSwd=98 z`v;T>3-e{>r3dcJBI`Fd_vYta_}zbe)s7xb4~?UkS)H4GF7j@*`Z-1P@`k%!y{8HI zC@U&us416na%zDgp{l;IVT*o|3UOczVr@^@uO%sDp|8f4{t&q^JdCNYKbwgZcdjT8 z&q#gMDy8oH-4#Z9c6%X?JiBsPc=ZU%qk)NjY?DJYEU3>un3pWQz`{91Wbf=u5Y z(#p;m6n#$cb@*?lviH@IT4q9|TlJeW+aEl37eDn;$r|mAlxW-y3bmZ5(rFYWbEUxl zx92XiTYa72eQFbDn_`fb@@~!~#5ldpDJqhE?&INvxAE(* zt9`p)o8SHQDTYJle5qeNxGtf!BM=?)b8vIv@NAs8zWjAQg`aYApRjKnPOU!HZNebb zd8(!7e3U6Oi=jgiW?8O&mxl+zGyl+?0TLO{%8DBlp+Gt?cIW%ABKo%R?4lw91*^rc z+Mo6H@pb}OZQc9Nu5Tr0WkHVLJoV^H+h6$jk=xr<2zPlKW&%QD363x#!nCFao6lL} zW)eM8%s+YKB1+y)KflGwiXX#>;}a8NDV0IfnVPI0kh zYUu%kf;LIXZB7d}b*25nET83Mr^Vt*&sAkx)q7sZS*r@2q_-lk6W1Fx$l&r&F8t_! zg^Reby;vrt6ed0_+M{i)7o0E&zES1VM_Q|jsLw@Hw3a-#XYG10aJO!$rUc9iRJ_)h4vPX(I=jB1)TDj!a&3pT5g zlW#dXMwg`SOUL=fvHJ9R}Rk})k^r>YupxA7Vi zTjjHov}NVP2d;UQY?(;5!}0m&@4p-$IDF&EVapkZT_ycJe@E3n;OC#64VANdyQYfA=NM={0mT`zmwwiYKZqLN zhwAV!WZSl;{7F}Dn1nug!GRYJ^@P4&FsE0Wfw4yUi}29Uo{5RP)5}IAY<7!}ETB%l ze+-(qUdM+jU!q7<>S{*=7qC(h6^A>uuJEeSoiL`SY*%a&4q=t!3JdH10FlA%ZITk0 z`Jf!#vz1u5HLBa}#=y8X(q=2kQ!~bBw`aObBEq-2hTJ|;RxG7skO=zzNQtIU5z(XJ zdC;6}Y;-uk*uBZmuZOi3@pF5l{8z;4@=`7XsVl40Cv_>tn*QNmI_B#w-#t`zuIVyL zN~)Na#=rmTut>n;yM9r7K)^yEi8$k+^KPY{^0M8!{EIZf_pkdBpK-!OH~-Yk$Ou$K zaK|#>2JY{@?dfsM&pE;zMLx(NM{8?3o5`%p$G!&!25D)nW#idea}9lePaIw#kwt9H8aGke-BR1(9ra9TfWYkjX9(xKDADfwC(fz zatm~*@o;BNb#-a*iC)pynVTl=R3985+oHKswmJrdc5ygq+8l&6_r4X-IL{PUd!Y%Y|><-hLe@ zd@b3csR3HG-^SyRdiVSq1(X!+objE0O5@>eHwBOj(M@lv78UjdQ5Jn7-nxIZy6wOV zyn~)p3Pec;z=VY}(|vE20~Zb3-4%j(50&qCCXznRmI&~X(ZX}g4GL1GnC@MX3~FdT zmsPOcY`H@&B!ut!ePIAhAZ?>m73T%s-XEs_D{{0~AlpG+EvEQw62I&F>4>DH+3H}8 zks3&AYx%N|sIs7Ut#>cXOD^(Wpns~d)%#x>lR+=a9A7E0cYS|Pe;ejQPzp;Ulehni z(7#1!U@X9?4um8@o*bwJviOWW2;OsLc}B*_t!?VAoc3ASr53d{sD2zDueC78mzm8?OeR20W2oWaK<|npmd((;EWTL9``LWb{Ef~J z{i5+)L_RT}(_>D~Z@ppP_0`l=71Pqs6vR!lsJ2={rp(PHsCvJTove_KB*qIiP4!;A z?%=}jxB5l)(8#WI08O41dM}@ReBcr#VW-pW&t4)IbpiuAWwk1vC;woOR9dd9voL62 zEtFMOerS*B=lfh!arnK&q-x^H_8TDeP1L{?m*D&RuYY|bYyCO0WIeSqt@G_Qjb35V z@!k^Y-;R9{deLYic$o&ro93ldq6W}Nc7FDKZj(6w?t!rwkwqOAhU8+D6kF~sWA<-#pmO9 zJ>O41u3!bFE-qRGlA2`G;Y%fx6~M+TnI-Lx)ulQJn3=C$y~^TrjfgWl zLrcn~y0(_zc53+DJI=ADh}FXyxeT59A44{gORwMb56%5ut(w#!|LlDZ0#Pa9dH}|i z?rF^1t$IbmF8qX;UG8h75^F|s-+ z6BF;Z4ytahTPpNb#W8+k?B?d3&+Hf!ToF=WmF8=1#wjAQ%u1lw)smx38=*JWSL7EP z;@c-*6C^<8tXBTUZyo312JS9y$zx9HZCwCiQ6gObHw)JQ?kdoJVUi^WBz4aH@NMk6 z2LC^{zB(?-?s?lr1wlz^X$e6AK}qS(T{@*!xN^+>^!l>zoSl_sF^8OL)ukkJ0}>*<>Ikf9=Z)gdj*dnbm)ulljOYXn123G93I}hYNk~8}HqQ=M zr`3F6OFgm(it6{OouX*zRX2uhpK%}BTILN(N%e!N(RWv*M1`Qt7s@J8xOj~jYHPMH z$ka&VbBBjzTwIJ$L~?Ra(pt}2ZEC7ilIaakUsebk5x!#1^?O*6;^IWh^B$Vj$#wXJ zzP0==!lXVaw2Xd%@l{_`^z8SoG&CBM_&Fp(=nI%0_?ZL|+sEtjXxH0khPKc&vOWly zUFa|csV(d{*>%p`ysh<8SwqUC0ln90Y~I&*y1&XmcsG?^-d$TyA?ZEenW(U^m~?y8 zhig0cC&+2ogt$&CQ$A>#0TCcjY26G6C;{!FOPki#uc`2F$FfA{BJ72}@HmSBJig zmIAA$>=eDYH5VC^g%45@Z};+y7*3DFTn2>L<~ydow!Kn7E7Hc1KOF9phU{NJj4P3h z@tyRWJ_Gg(iVRQ~hhyu7yY1X(<>W6HcRom_xrr@tFPeIId|&7!bd<6OF&u%LA zBkE)ufE#Ezuk-F4C?RNkik>MvdI*}Y(Snz`{|#uW0b}KxH(5y8GJxj_+B=RPL`Szc zIy$4;Jzv=V1Yv28k53C89>%?gkMQqGV39-?+RXS9bC?8^%&#thi=9jVDHgbvLtw8B z01b>*g-c1*bd8E~@!MWn0UbQ$aZKKy(Zk<#=dA5e~k zU&n#&D=-E&9CB)O$AZn#wfBiWO&{*y9pq}eJQfgu0Wq%O;kl*&Ls(}V9zLh}-co4^ zY-nTrI2v`4wGVWAu@@G=!BS3sEF7OhhHA+RrY{hRH)8tr9bYa{oaz8n;i|!PNpnN(u0R@pR(rD zGBTR)n|UcH?AR$6#bS|^z9~(gZ|_N^nV0}Bq)1jZA@qw@rjk#4TZBxe{PQntI=nyH zmbieR$PoFM$P3BAjgiM7!XxW!Km|Cz}ifi z($vkK!{tJ5cOjisN4o4XuTHnmAzbJM5QdI?KqEVEM;~u#W+WP%$akP5cW-1D(1ao+ zWovXavy;&M(k&+^9H$3i9;EuFZ8O_9+%G;%PfL`R3aUVJW4StC=^&ph*RivLtnc@pHoE!Tk%^Yno z8RzGkfL7Wtl#w>LqkFb}!*M4t!ll&9-MHfwCY{FhDG?Hmjz(vjQ8Am^mCDqF5n&q7 z4L_YAG|n`2$jNb2SRo)KgL+-&m+N_eYHvL(q|(pN)ngZ2jyy~8WG$U>Vou(CAF(`} z+oK%a4iZM{>rFQ&d(+bTn*t;%s!LjPLFV_87a=XLDmd0SI{Fz@fWw#(f+Oh7uGQKd zkZ;W6-f@Yt75Svd`^8?gbxcC6fPx;;t4(1I*tGx~TL1sJrs_R`wBIH47UEA@g5a<*wb!=Y{N0Z9X(y~ABX<6im#ZP`U&295&0L(zfmBq-(1;OEE z3kw#mo#rMMYrOlNMlE$b?!)cl32g?9L4m=m`0##9_JDxJn;0d_<_W>1JbV4YAHU;uDH90qwsV>9)K#29w(9(?~UqOF5ujp&RfCD9}yjh9Tf zaNU0MuG&q-%*;&&)5Q-s8c^A9)hCfk2Q~Dn(-9#y!Mw3j)Z};qYCc?vJxZE9FAogW zlI7*LCPB(Qj7y%#%4Hs%0$+fHGt>_ZA*AK@DfwcnH0a~=to2_}uoC#S7NVLA*mnT1 z*flm@^PFZDI5fXT--wC(cWwX_gGm_zOzCn{l9>sBMv4+~kq*?1%k;JHDtP!O`iF-V zska@S(G!L5u@*uY6fU3|Qg*9zZS7ybzOuI7-xd!fKds2$aMIA|4y7`3eepPI+kWAz zkx8~hAl7WHzdvSWWe%(p*RN^CAcT+fWoUF@AV5&(80NC%uCY`!-+zN=XI%vcwi;QGQZl(*CXrp2+MtelvU*OvvG2X}NW0 z$C~S@3VWX?3o|Au|0oax(MhF z4d9le+0;%~cG+K1`5!5va+zhEOkmv?-gv*dIxsv9rKM(*rp~y8ZoW%a zXG!*Ao!f2u&Wdx-4{ZHey>oKqLO zih_cvkpmN5I?GF!ubG*CVGmx4Bh_WGJKIKGofH+po>lEY3wVOyTyo!L;q69kN=p1s z@;%gbJoG!`fUyE49(XbpF>BbXR~NKwa!DM+Eewo2C#z@fdA@??78Y&o)acFG1ruv4 z`2$gLy)pPlzHaSgo6;h`hr9|S*zl5GfJHI2=|bb_4G zZ2;Ajx0;easJGqJ#c&;A4d_&DHNuk!ROfcKQik>eSCj|FneW}BWAHgsD@OE*2OgDX zENmbleb=wQB*Pv_(p+0wD04q?V2^1%m{Xc7GW5MS+b8q1EiJ`LZ*=q|tBlL}+;zvE zS4IXdJLu9cOhw#wb#M^G>JucyE;Cy*`4VyjA5RapnDOV3XUjRVtJoR-y03^S@I_66 z6(7a^+*IiLuwkMj^5jV;no?2^qG@W-yFOt#MYe;*t&lMM6o&m$Z;+6%?mC85NU*$4 zaT=Hbqu;YwzKH(%S_Wmbm96w{dJ$@vcx62OTCO_9vbYl@d!8Y0(SN+KT6TYIt7Vd(`GhC$6>Rszjq zZib>iD%aNXPX{N*-P6D5YrI)n87XnJ-jUW#r{_M{B$= zYmC>$@&C&3DLs#bsgTlM?OWSB7-V5Y@zn{jQ^t)Oc}hgp^7Hg*!2lkEtjk!rYhy|E>biubADi?>#a z_ebDujlK$$FN!A8GfHd)Q zF3{;jPWeD6IE?Hld$}*fYNU5Wx!(1HC@3iNa6SI1#jvQJo*A7U4nd;7F0nhKprG|^ zGr%)DDFp?>J-Y^8zor{2^H76|>gJ45M<=ucEM1R+M^u-`b7R~Xk{Ql|b5G~Fy$MjV zH8xx6499g;IWTv01a*3lm{2DxA_Q3X!|pEoB%j~I%DSdNPdYpA$io5Axz50 zOGjg<{{?OEtj#c1maDdP_#GB@1VCOAb8`<^w8ZJJz||PaEJ+0|V&U5Id0n*LZ8M?I z=DKlS7pnJjJcVblj|3{E_S0qgZ+O?gbp2`ry5C>F#_X5i_jF$iByRHaGg1|_eACTu zX4hUOA#qw=oyn#%x^9Ptfm0_X1)2#guv>JIwBBD%X*H=xL+_OFGqtx2I5u^_z1H-& zLTxWfv$Lt6KX>lSPyqbJOx!VMQW`AIyw%)aD-~8&4)S!&sUw~}#JcW!i)qe|jUh7~ zGzzh?4z9g)X#w3rG&Ciwd}pfV1neHfE2Xbrr`FWO9c~QpZQbhjsjfjgsIjr#)*6Z0m5M@(4vT*z^Mc2u+24aeN9RlzFhg>7xt(ze!6Gsq+m z%uE0%YB@Nn;s=r5;^|6vvP@;=3aNrqJOI${fhMzfu1?g*$Ygu_IN&j3E$8I0|5~)| zWOsM}P$5>Z-f2F-P?Mp zkXJ-7fe#&Vasm>t#8PKxG?F~8q;KL_xrbROM;SuHAiHH_R}7&9fJl_GRItP6h(6F( z))^KSRP8L7N-EuXU}RK7PL2UMM&R(UuxA|4ii!#!T{?HayE;cz9mB@Jpm~nV2C)6P zMFtl2`#`D604BjjL(07Lo%Z^w1&q64!ec*I1F<%iSt}bs<|u=*ws!7^4oH}YsSYaa zB?>uSu9h}5+^MU>x+TICmK!UO$d|ccQbkoQnz1pGo2yqf9R9@1{K2(owsCT2f_2v&BOi(arn>Js+HL_D)#j`)M%(cFkXrRTx z?w}SKo=JTf6{l8OYSeJZsW!z9Z99>m40t0fzZ5jva=O1Pq9gW-x($Wu`){}nWGG>P z@CG;8CcI#CV#jV`s;nLY%~Z(g8*>e=zqElmq@TlKwn zY+_+?R^wvc*jUb{>uv-twg^z?^E#URW(}WOS;5ZNpjp^;;Y$KRJV0qCSPRaocl(f# z8N+5E@xTnwP!9z+|7^e>TQ75bjDW!|p-Md?bgLGh_KHEgZEIU!=h3E35YPlzR{uvn z@@s(rzY)O4QVKQ`1$8p0y#xS$j3y-oy}31S%*?C@tX_fcui^U-l41H^B2fzU=J1oG z^wlsYR!;ZEnQl~YTeQ(BEGURAKHdvvNnOi7vPOK4i*D1evn^WIiv{gWop-U$D!jaq zqTz-5S4&y6VDa_5mYtml$`|@3Mn4v}uid&=rJ|W#Eyt}Cf}Ro-ZFLLJH_n!%rJlO^ zKPuONDh~`qOHtzjPL=VE<9vVc=e`32)oOVL_Lb5!sh~RB<5ufisi!L-_+nU{DRFJl z86XJG-d>NunFijlWY{81@yME>s`ZpIkKGY4$NZ090d!D+fMr0DnwZ03Dd?vGFdUjq z6;RP7qNDwM;~E6M_SS&9z5nnj8YD369C(B*N1*xlFuUpU)u@(cJ|Q8ty`yol`dymT zU0)o-q(>9tVhT}pd;Aa|;Ta&3`1S9ZgezC3gA@S^*q$fQcY@d{hs!dN&n&yp2~wt` z+>$`)4!D*0HX~&}S^>VZ`6>5Pg*luBy;|LkzLS%hAks=c3SSiY=Im}wvXC?B4JuP> zOa7R|JXXlfYeM?^`dV7?2Qhz5d0n}(hxMnOZq*xu*wmEX?s)6FHya?Lw;m8|)yk+% zpGS%Kj<2sDyI%kV$(l?szxv(|mBc`=UCCF6l!5}lZ>UY(rd=xS!~MKhq;*ls91K3l z?(gmbA&;>C{WBW8&LpJ7E>#uocGFkDY6$|Q33l`5G8k@n93?SW(84HEwE(L1;G2Z$ zXLsMqg9M;Hw0q>1P=9iJe=dt0oCjzO{9@kslKLCLc!mD6UO1F|!!b#$5z>k4Yb1vo zVjPDTi%N2G3)w0b^bzPqi9piK-jOTEp}RFjMFuU{ev-h_KJoBeXlo#=?|sNhRmlJ> z9qxw{&ZByNbJtzz=J?Yhm5PC5=d>P9LNdx2_1x6-F=XoU>+|P{s5au|ndDGZDubMv zxoOYDtz-SZ!Esrg+oOYnU<}Pke;RXKg=qO@96J!%yL;pg7@2^)h53JqUGiOXu?I09B1C{a z3Pnd91}NYd5bi1GbbnLMUEvq`4ZR~#P-2}7)%D@`EA#cZAmT$Um}t)u{qm={Dg={V zt)zV#83)UTmL?`Z`^nsVY&FuTmHxmK~i^ z9Zy)Yb8<@SY5(4mFR(wL5-fxwv7-wF`-IKSV0SlP*oz+$ZQTCwNOChEL$>rew{vmn zbCt(nCjaYy|9KJob9#!IQ9W(vieiwyD{irOr`&M zzE#Z()_4JOW2O&}S}}0t;AVho;?06=(yRMGF>-J)mL{hgoe@&4*zv&SN^!PpWCy|XYLZ1n11SfY%h+)3{yAQ z#=$;?P7+P9J5JWP{94>UpGk^8YcN0C^2Em{e{c}D=xx_<5x`p$Q-Ttsw_vN2m zIg0r6_Ld4s2~U_vj9_5Zh3?OJ`)A0i#KN1z#dkgm|GV{@H-10O5Rm_c1q3K%P~ZOG z{-!4qNq@KLPx_aU4qfam!x$7??>+oBzbk=FdSI>+_0Pfq-x2ilr>Dos!@D<5hw>sG zlH9_;Nd5mGLJ%-q>Zw4FDk}b1(U%XvF8_R{75>aeAd~oc^)!}OSN#83yr`{XIN$-K z4G#?s{ne%V9}`u`kv+fp+iH9tfFrT6t!!u@ihJj8qU|eJw&Z^;q5G92K|wv0`PG&7 z+3jmJwUkiRm5T;)q5C(MJ;1_KvvS+c2#2YY5fl6U`&jG0PTkQJ(4$nB@remiIsV`} z#w4$!^gq7J3eE$RdH=uF`i3#x6o6S;`WO;6gJPvYkkF_6uS8I?+*U4EB7mk8FhL+j03J_tjJO2 z#AY>1W5bq_PhE(ANN=t9~-c>Uc`vd8TGxa?xE1i z3!CcOrjcqkU;KP_o^i5Q%i~8I7a?PJ-hJ7{+g9$92>Bo_vHrm zYK}?rU7^sa-b=U7Ur~8h!0N0gx0&x_wOBlzk8;-&>jQdJdGh6T(U-Lt3Vq3!DW_bQ zn5n0en*>ypX@X>elZ&Y#Z)HtDzHacZscrrmdH7>$`}?3(`Qn^?-vM)g*HcxdlzkDU z>*;tlRC1F$8(aHDQje<*Oia_xR#+H7^(}a~^66<{km^)SzRQP3qH4iFcd}_*TduAP zRUL`&nqTSxV|O~cojs{jfcvU8OOkmC@wiS!6BM!iv=~DA80%CHm7u9aZ;xCDj$6*o&LV;&IGfk^o%||(w>}U3+B7nX`L;=78F2;%nKjP24F~f zRy%AR4_t@luX|mr3-D;53t+u%Z#@msyPJnKii-=Zpp9-^+RK(V@q%N$X}!PeY#1Kn ze4#(Okx@Ez%8*hOGr8R-7=F4uSL)zAZH=DboR~Sk-jmSh+!4H34w;N4*K~2^nBP7e z(^J@)CDn=8IXVe$;dTlYq|_bTp1zxi;-%1*UhNoK#;KbefCp4lBuwp2pKZA5tyMWX zE#`4lFRj%%y0S$L25LRgNHm>v&RPDJmvLTBayVi3q8nbdrsen&YKAkFEnt_!bErxX zMVg7NC+KqIx*Vp~R;?SZFfaOjA5TK6j@NTyfeQoeMpe3Drj)gk3I(T zF6ij;(A1mCzu+dTeGG-W{`^=atMh-{&%{4xBN^G|-2I5vcefx4nW>f4WTqDd37E4o zsz35XuVZS(p`uM_eZ@Ps?@+9BcB`1O5Iwtn{W@0!V}xXag@t$p4$1n=^X5M0)_@;w zxrd?CuNy!3D#M@3k|Ww0&p}^`LucT-cZEu+3qhpZf-fq}31x5V@)`eZ*GOiD$>|Ua zxELR$McLEm2gVk%QyVqagc9^nF{en<*GO;q`^*{=-DJp(PbhQL93AZb`VpxqrH~L; zX(-;6q%pD3ouuKSF_T`oD{W6R85-fT+rG09GtYK(j?a?&I53|{b0=6sTcxHV*U2Co zoF0{m(O~c7yWE9VyR#auLCl+xxzb^u7kbahBC(i`)yT!6TgzzAUyP>e$L4;Z`U8H0 zW8upso2g@tPh#@B9t8ER-4DVb5~A{VJ}0n;xBuuVX$el?>h;om^^Nn=3YX;m)1D!< z^Ov^yu6W`8 z3L$X51rQO%#e3A&ytq__yJq-l`nx`L0AdOgX-?$`S5o$)t) z;)jLASX=>3KX;Gm?fiOWO)QPq^b2eXxl2o)m6b+={qB&xCl}ac2Yk3tO#VQlOt{l{cc0}I-9#*F4qgM_r<+&`L?&CD6gi#U&KAC76p4UartUuRjf zoO*(N5WSH<=-Egnem1dfdz)ls9PG?IoNj&_X06-1dB}}pwDBVp`oR; z^m4y6``?i7g7(+m)go4=o&11ax*IL_jvp|kkP3K!n)ZHqj(dD?K8*(s?ITWl%Px-ekK z-4v`_Z#UO2b`d~*M%~~p!hi$k>oM!-^U zRWf^6FX~A>ZneB?^Dul{^}6^i3l{-bJbGr!^=|oMmbu78C?5T^ef2%~9oz?c+y+e4 z>^=@OF@ViRTt{Sqlz)3pbs?jQilc9`IkJY_VRlVlTFx-Bf`@550lyzQnLWz~l9=DdjGpn1bN`{NQYjIRWqYaH8`up;h3m54jqiu}AT;xRUDHyhFjElQ`lQUW> ztMkl*=iki4+2?oWw8{dht-bUC*Q#`~ao>-1QsMAH1$9z9eBuR4VlD&E;n3sPNwm8r zuAp5M@p2@{XjHUZYF3*eI2{*vj>jr;b-1hIMvw#X`ev^D^lg|IW#(%UvQNFq7eCjh zpAQcNn%(!mR;b3TOXl8A=S6x)XRm>iv8B2U&SOk@Jbe=t=E-!gDz%RZo zQuf^5wyd8dC4ZaDh~r1Ky`&?DMe7yzh1By)XTmBI*>|WMj&yQ6T#+Tb>?kQioy*mb&dvJ!d) z=$5OiFPxv`zH?YW?-Z|BH+Hbq>@>kwB{Y$*{p957@OnTW?&DxWQ>lgq5 zKS8Tx$AQ|5z9zNWlA9b`17gZ|LtEWRWi!TPO!;O6k@=2&LdE;QQ?_2>v>Y5 z#Ry|EG?8~&+#S|tjcJ+n`*_VwOYDeqGraM^P`&!)j=UMHT!0xT>UHRQpo$vRM)wm297I)=GIAzJUHk>(Pu*ky6+0Wz{?tueNaF+(V_RFL! zRa9FR=}a&=SMFn~nG3-KjT43xGs{{P&Gv;|Fe@f1DEQ$2_`6@j3e(rOcIE*i<++E) zWe@wltbGee7^|x(riPPZh^;J58!}dJSB~3EvP_|1vr4Jdw|9~(m@MnIaD)rrZcs%vZsTzdZ!PkIm4Hl3^kIKQ#h|KbxbX zH8zCsBza*V3WICv=AU97)(apw9kfwu6Ec?z7H4U;-oq9v%J5IPJ%MdWLjqb;o@{dE z<*)AFaG6S}LR>7Jt`CR-F}#ZQG+Y|wBffq&?Jw1`^jx^4R)L@2LLsT=%ZCE#=J1ux zWk$jPZp+bMul(mT{X5%VI)_?X#t59s>SUZ$R!+QhQ3lDrmxn;k>XE==+hKn()K9Wu z#^$E`q@qH6b8%C}CEbTdmw8_&G?;36G9p*57K#^SZ-}=(;qMiGd(NR{f3m3&#R~+Gb+~w7m9mbw$6_8yx zd$dtY@EWqlFyr2B4str)7k~+w_Vkbhx-p10U`kf9yEX5k}#S7283V zoNQTzc)ux86^*?9e_S8$pU^5hZ+)lIepWnqtp>J&gjChrT&{w&c*xpmO+K!JuAqH} zsWY$_m%hFxB6~~$yy2M|NcWshe=hjU2v`uJ(e~dC#4vyL#|v+ZNckPWrcXiLr+wNE zRQEp`Yy(2dU*(=ZpFi%F^4-A!YOyPb^7#rSE2RFC&vY`E>U{qpAH|Sc;eJ%fszdx8 zPyoDC9o@{(;UVY#(HGqKvx9=)^-CgoNDD76WkK)xLw6id(|*a96a|MTC(m4*pZ=dW z`|Z?!{fc$nnVqoyw4g>twm?puVtJc&Z0%L=Cl4@bVD~(18RW{NfuSCw=ST7)Mcf)_ zsgwPZ9lmBV#J6b+7)xu(1Z*(O)Kh>{QoCbko*v#f&lD7tt_*t6G5*&pS=XJI?F<)P z8o^I?EL1$coB1F=b}ndY#m*UjnzaID@x9=zjm%4j=eBl*yb*2D#3VZve&*-Fb2C>Qu|v zh45}%Fybrb?`3I1HbFtg zcUI|ZDu7M>{zB*7t>!YZ2O|6XBYdpfM#F^+_wH38K5mKtqs_?3_%|b%yz@JQ1N;mF z0YsTTAI~C;SHNhGUr)~oOK_BrEjS-h@;L7Sn1BG~6Vv&dyO*1ndtl0Y@9-!IsmTnH z`GsNpEi88B4*()?g=wHvCMH86;Rqdryk!v_Uox_9zsl=5iqsPO^}f!1o^0zm83P`t zE4!67OQE4eot?iQkJFe5fd=|NS=rgW7k}$m{q^d=b%C!%#lYkfPjF(8ca(a_-an#k zT`!Mktm3}`%IaP`*{93StPGz;g#{Es$;k71j~}`#^p?t&e*K!8mpAbLujSJDMHKr! z4B*<-$)Yv1y_o_?+9vFpYb(JusI+0nE*X3*74W2z1duI|ptt959&#V<9)J}@~6ip_na zQ-6!9{SCHJma^~1TbszmH_zCdO-yWF(ipF>si-hML%kJ#{iEYGo7HCbI@Pbcu+PPn zneLgVcD*N-Px1^H)2O2K-jwpp%Rg{%QjI|ZL6cwlYybc29$m?TPU+V2n4dUc$B<6X z?}NF!%aLlF0{Y%2k7=)O=M$$B)pwB*lecSS9^Aa~;K2=O?yy@^@aLOX#qb%gU3DwS z8oq-0xOdG`WzhewbO})?u4lxJ2Y6!61;zIfXqVs?V_zfo-L8(8F77bSddtfh7a-(@ zU-PcXB!px|IS7_FSFnZ*`w3#2r zJCN9v)&=j~O^lx0)zF1Q&uM2=l-4l#sE&pt_xnC(b-(A~!T-=yTvV9F{J8t~Pevih zurq-;4h{vMV!h)aZNLgSPq}eAZUdg;O8CwIrHKRFQ+ zk?M;W7OY*+A{9L!Vvbm8H=Z*&TZk4Nm7gRmpje;EOT|uw?ax1;F9L78bR@t|#l^v$j54aT*h|VBCS?d%?l& zVKtwqvU7L8DYT5xj#6bjcxB;@csic!hb1%XDc$ydaJvH5{A|5{km2doUw7P_`V97{ z#`zOBxB8bJn!95Mlu0WeAQD4t3-=@Q4|<2hp2b<-*trhzxPnQh_8#2pPp@M#wN_i_ezC}3=80QD1O&y|_t@&3v zZ^os(HCf?21nQO~E4!>FoobO@B&e@bp<^XS;%Nmj{5}njoP~&ZYhj0V?#zZR#>#T0 z*D9aqFR)CO8}HOC>rNzgJj4?B4R-kO=!J#dX3^=034Jp?W=+(MMLdXfeABgkUJ|kypI8Pn8T|gTob98>+Z~SqmK`!#&Q?%!bb7qUu&s!>1G#)FyoNL{Jy zFP0nagV&Pk=w0W=O?9v0x!#gpd)YhH)PkDqY$}$4;hD>b7~);WGrWG?1DReqQ}JM; zAWrAAZ#+RVogjB*6~Kub$zRh8yB_201c6}?~Z*{joYh_Ts0e!6>p z3OIi3`-_Vn+bic|87r%P8Z3mqjmt*Hy*ZhQ0jZUx>F<+1qUUWMYjix_SP4KzJIu>!T zb41&{7!M{lougmxet&Yh$*IHiU9nEkw{xbjtv2s=d=uf0iy!l+*3viL@i?_XO>eBjP1XFJqaox?eH zd#)uTgL7uD=ar!d{|9|IT=LK=N!ga|^g99rwMnjE@=UbqdeS;Z>C3Pzl*6ww=o;dHLz{8 zu;@H}71Pj-XWx5xi<56C$!fE)?_;0uVjM{T{9}rmVaZ6X>nPW`US*170BV@@{?lX_ zw*<#6?`Hj{%UIEM(51mn_XVTFM&8WGDZgxHHH-9&{%AUZ8WLOm)5t>o`U`=?V{Hm= z`DDddD2Y&dFA~}mVM~nFG!}NlJwGw14x$#I%WrYw;!Hd~WGgAi;3X*bY0sYQ<*>_T z^c}nPsscN$*4@)*jkjJzxIrFIE5kdGBA+-D1Pfe9END!1f5})P^{D>78o!ABijcn| zv{~yNSECUK1aM?Po6&cW5rQ^Gj67&t>|N)TqRqQHe5CujuaZ-d2LNuFzAkqhU{|Av<0=g z#WZxmx)7w(`)(K{zKu?{>^KRNZXGsia?8}4Jss@W$s4hyVhmA zhmMam3K%F|uF;gVb*4Oy7Zl`dMMM(jtFxnren`r0(BJiuPvsaUD|-*sx60leAL!_t zXlU#lZGF|C|8#`qp!-5z!dan{hx5?DWGitu8K;1i2B*APJK%90#rt@R%dJ-C2&PRj zoV7mq_fiBIL4N)QPzC@By~6MHKh=9#de;5E zEgToS5=P4W(Nt@>O=4;>rhSPzVXNaJ*RT1iCz`c7vva_R_ z3{Llo>!YhH%J{a;B}8>xe(u*a9VG6x>QQ+KJ)Y(nCdOJ?EqoJvI2s$KR!kH%73})f zaU{@>-@ze^5BpHPJ@C__IetOL=aU>NL7KT&O!asJXjB}^_;9=U9KG3ZdU3Iy0;U@) zMSQVU>V^NdD zRV(^vdZ&2y6t-Y|L!h20@1T8uo_)u;z=*9P^!cUPLBqT42~&q#I)2}wd|44cW^MC2 z(z|pzfjT}lcvc0b^3Dppn02@t_Jd>N%h(pgHb5m@c>qE>F~p5 z_7lX)(QJ7&cw+e_3*m^~TvgigSNjo14BTqt^<+vaxCM1{Q}-&bfz_+z(du57rMbQx zS+pTch0wQO?A$YTJp?I$cMj^--m-A&gKWF65$T=={aZOO!aAq`smiW#o1TjUzd`#$ zA_z?~gs6{k)-^#k@pwt|NA=Gm*CtL?F~gIS*K5T$>j3+Ld;XW{VRtg^JTe!t-Uuyb zzA7#Mk2Q>j6BqlGD>iylPTL(&))-!(hqKT8PYUz(S5L#c19RPsTOq`KJRDgockR#? z+MO_tryQk@$|sBMQikI^D~f5jR#5dzh{QEEqbD_Q&3cd`_Sja48b({_Uhj7fdY88j z+}zJ1pTci|pw5|nyx!QrL#(&tvGf%}3+UTbGuj3!u$fchppH3e zZXFb*_7Pif$+RxsJtEsle?jJSXhpa>E1tynG-qo(97H>2(or1Rf^}WhKL>aRgCS8! zGY*J33O7gkr@1s<-7pK-5T*^$`LvtO5+^v6bK~gh2t-I{>}&oMVGQyK{(bZ=aBIv~ zG~;}8GW=D9wlrRT*#!9oiZu<%>kuOJI^NTwp&}2C$;;0IrFGf|TV|r&vYXgEvR4tj z9tI(`@$}m_{n;Haautl3;2=z4!0B*uEQ^-D!SX>Y6Hl1RT2b#bO{EcptxT1bDJP~M zMgt1 zetk~F@(AAdIP;?zYPC8G=6Povj1)kFv(gz zelhBvley2$R)PUU~KnV%ttfC?DTXP<&UJH#2*`2ddE)BXSWHG1SqWc_5r-A4e@uSgS{a~G zKrxsvZ?Zt|>Q3Q%7RaUAJGYcf=Ovo$88D7LWJffn?hhvXLa)k$bqwEJ)n`uciN(4V z@tw9p+iH4ma$k?pqn=}vqCnuQVVqwWMPk>QQF{FoQ;UI!kF%#HJ2$U!1&|hH)i0FD z%t!gV?JkfdlCl?A1P4?I=FiOqvRTKTU}b#*RYc zrN~;Ihx&F*V9sPTpSZfZCJT9tiotot`yeVAw{G3a{!iefr=v5`6hP$W`W6uLZXqFB z1+sEhpj??Ed)|ih#6{F`ziTomRA&Km*G7+vio%_WIBc9tu^&#^_asa*NAGD7asQnA zJ|d<@qbnGwf{&m)l(aN8=<0bT01EYX48eBUt~<-D6^_P!UUQ$0_KK4^Zez!<|9sS%`M8rQ2H?~)9%}+R6C+YmuXP&XCpx|F)$LU&LFp^L0X?xDe4LJ+_ zPMgcJk1*vUIyi+u;%p9noF(O};xj?03G?of9)SSEm*Gx}vdPZzHP(325{E4{iPMs_ zxZ)r2H>tv?&Oc#0tQj%Vao_e9Qhdx=ZbgsO%B-oOOub;%YX@Pi;U^wsd;czjP5jhg zaJKAoj^}{E>*I}+GH>?h+~r8wd)U~j|BVSoM@K(=_^@GRm8YGNq*3w-ZQxiI##c$~ zSGSnC(AhYAZ##9_;R9;EqEfvrMLcaiDLS6(&_uU20OQ#v;s>*Idk9_D=%%WmA6gRL z^%@WNne|{A1poo)E|oob&5Va%-SYl!glj;>W>++IY9Izf`WMS4o+2*0UyyZy0-Bj% zy;9z&R$6+;wEfL%6oRG$TY_vK(XmTyi`%V4gaJ3iZ{gmGnBIPo>dPZoos;!)K1xc# zD*xg7rf6!iUb8oMvp-kFXU)Ev2T2My~QO5i);N3zB6R}@=}-f-qg z>{_&!IH`WTA%DhFnN2HTJZnT4fO>5Yf=S-reYKDHMnIFD zE$5S};@o4l!D#v<0nWqnsbvc;=e5F(PE>!D-AKS`Pk}$lq1%4MM?(ygJ2!)P=v_jKlE)p*f6$#0^1Hq8D28J z;KBc-<a+xosZeY)O^v9H`3^i<&PSv2Y`8JRhTt5X3 z17Tym$8d}`>AEQ zOcm$$em^UIkD@JvoSisY8F2<;ShStnQoy7{^1N&+r6g5~RF^QiYg^Z<%Bb~xB77Nd z;Sut=uWY|NYZ6bpLq&6RKysx}uK|beA zrYnAisEBQT3 zsFN(y>WuE|+`{_i$E&nr^4lzjVAS1T_9snSYUoB zbz~K*P)i|6)tSv6WMXYeNp-bHqdZp&Bh%M@cqhsSe#+4|Ae5TofSCAi?jA3t(jeAq zN`E`YQ<53YtE)-@j$-(3^3wdABt$B~Z~ILiF84MArX*ju9>J@K_b<$kq!1DPo>Kes zHI(gmbZv(;-k--iZH24fS>W0_gs!b_vnI0?O^V(HZL#$Fbel`MQkG^O`6L^gIDJkk zJw0%Y%w&!eLqlsepH6u1>zF4+rlZ@kD9>t*Ep){3=%e)=0{6|Ux*TExLVPeo?7!Nw zk58$vU|_VbUKR2t2lxO}a-)9e-y3F&zbMue7E{M-@s@W^%cpmd(DY?@ z75Z-&^=fd;Xpp?K;V}0h5E9S^UT%^CC&=KjAIIB6*mGntmE(r1ySw6pxIb%$r!WkI zOzC@z46h9VRo^mgsVSbojuq-|MKTE@6rq4Kud5oHvh@HERh$lhk$*NdDNT*Ls4|4b zk_9y{S*_~OL)%x6YmRfzJe_HmeUIpxSPcS|NDo(;YBAq?qwU={ht|2JGzH{>JOQ0&~DY>ZElr>u!VTB$*E2tHkdd`HjBTs-gdHA4wR*SZa zws*-o+m%QVl*%RYfxq3^LT8ujOp`d_b{dnWoAIH~kkfJ8O9^%6fM#pWck@ws1Bm?< zu|{Om`K|n&T)RUg^3c)Jkgtvf&UfLL*x>~RQWMh4T#hON@o?AzG>ypNMl>D~CiOwo zwb4KW`>-HTsO#dsz)(AvNaf&VSnpZp%}t#xXUcnD9Iq?-;^-U$tvBB5CSOCdT@vd& zt{ngCZkxFq7Gz~LiN$bOY5$DLw~|hsZqR=!6+b*Szvv)WWJu_J{bf|8(=rXw!LI7_ z$g2Us?1*uqg8xL_k=t|}Rfa3Hi|7gQ0~=&?c)oCI6zgIdRK&vC{MDU9g1&ZXZNhpd zlYAp&>SrGf&A-(rHm`cvxSK2m0kw*3?$m|2^IE*B)HN9X831`23G43gDdWJAZk|UX zWhY8b+fdh`tG6TU?HEbY^fHrZO$l*9?AAEXwM;J{M|=M8y*UL^+88&n!TJ`lT{%*p z_#7K4lm3SQ4RW7F1`LF1OEfq3FqPi(B1zX{3h63al8*%mMBD4;P|>~R^|n>uR2}EK zZOI!~v601>qZfFj(EHCuS>d(k?mpd|noD|XP&Q?bU)tY;j`%6~fWJa-J9Y#!NcMl?a>NaH`)!jd>ZvrO8#I#Vy5|6+qu=>SL(W3I3)~`#`mw zlM#Hr&^Hp*RVG8SOPd~lcZ!?Sm?;%u#^b+Sx3}dd&7}{_@m+E*OWt@oc6#%AW&^;3 ztql|TSrMrJUI7N#k$2eQ3W-$DpMd$(|0Bg#SH}gK?Z5GyB?zh( zabK^H=|b0Fd-vdLDw2iB>x>NV!I{>}gZ_KAH%g@nP!c|mvy=WRk=*vo$xj)T=>xeQ zIylWnD-R)`s8J6bty0c;ojQuC0{PRJ8uJPY1nEFxa{>|)elxcrjT1c z0f8*IXHAUMT@-HR3#l*lUeT~7TwjAL#!HqtUx5tk+MO0pq)LPBX#dFcRA8ggAD)t) z<$kP<=`&Qrz-TZdH`oS1uKya#(t-jiF!5UzvAuj%_|6;nx_oJ`5@; z5$X+B9aVu^EA%zg;5Rba-+MzrrAij~CQA-ekxD^FhQ%$&s4H}>U@LDQuVYu6f2du4 zVOSCiZ+4@P7<%}biWpP|{M;Q^mpvW2O>!#@0zVqX(`^tG5%!6NjXEff`-Pzuq9o$$Oppo~+qyBN<0Z$@V;bm7IH%c{*?Sqmh zTcUY#;Ocz$p(7w%i)q-Du=-;9T#ki{5cNpaJg&A0maLHX3y1(pm$DQ*U{KuUt3M?c z^*bot>nkNdy^)t!ZJT(xyxMi{@&=V_Xs&idBlHO20^*|5z-4h(;q9^0x_;(i!>?W} zblneP(fQ^E8g`5M(1wW2!XC8VqrESJRE9S^t94^+Xv&1F$Eb9Iq}whiBT<^FLNZsb1%fZu<#zfjSFW#|<{Ldz3IY{x zmJ2f{&5Y3SX4xp;w+(I{5JGzV%_BqmB~zg%<8N(0HrME`4lM8lMSHoQtrM!850qz! zI!YHe-0X6GUZtL`eQs%_XKX;2%0Bl~xg{!BCOhuZcFMO$LKMjc^=a*T?sDiM@#y@G zpozfLZfx~4fquO7^d_hE7K;%ii}W;z)%Zp9aTfpy+RU=2M^E#x%Dl7<>pwh(U=DUH z9p!)2nSIhJBW2;T95|}Rf))J^^2CQBR*VWDoc=~A=t9nW6PdSF{vHc@@1F55# zVs4Unj$0h9St=9sBJkyq-8utD%9|UU1@zK!bK(g6Or?+O6r!=*>EWmUA z8l4l<`u^(7X)r}>nrB}^*(h5QpS!oMW5jcrqPphOeXmtoY!EwSMKPr5&lLUc;mtxqvn$cBBW|9uT{ePNx%v@2||E4M1z1 zD6b3GzL>5hi`pgbKe=!uWsgI@g*4

  • Nb0UoqLHl&eXi1q**wB|{OvfaoRVwXj9y zGAR*7uCA^94-ZyJX$;te;-VrNL`39L^)c(clwY-w0qK-aNu0rpwjd8sB9J)EWdhl~ zf;G?tg$5sxKZ#2diojr1gcW};&1HcUS+AuYFTJYHLJc>&89$<3gY)pif93LV*}Kei zK>J30yT^cjPqf+n`Ei>#gr9FsUk(tkCt5qKyIO&Hk^)7hki^S7{ zG@4BnUx;CcfQaHed@Wj}UY-&mQ1vSay)L_zYxQV-Dyl&>@RJ(XS3OZJv|s}g6kb!% z2k-K##%?3*zu&}~p6FWO(4&|U(flq>X6}sF7UiM?VZ?DCmEj5hxT>Qc{8xT~%)Gl(=}YLM|Vt ztNnauBnqR~l#eGAi}&F$+wrybNjK>O;A{THXLX7?Q1=Y_==nproDxk(*s8lnenOir z=ChB#_yT@cm5eEW74G)I(*f8v=53BN8=DckrTIrDCC}5Htk!@$1S}eqLy(k>@1)ju zqVsb({hB%S1fKcEcP6hpoz#rZpFqs_$l1l(Am6t*qd>ATJ*mKj%$!&*HXrp_8z442 zT$*Kv0m+7(Uy-*@CuDdfoKYsAyVS6i5UWU74Q zlDx1uq<2Dcy#@+tRpFh~QEx2)!w)1L66t(SkOxu*chWT+bJkX9Bz2rAyD30-Hk&K; zx1x^?4IOTcH`nIm6oh{p^mrrm0MoEW92=H8v>1iQ558m5te#Cf|Wc z)r&{!BQO zbEDJXoZ@73&IC z8Yw_z0wR-L_5r*fPag3Au0GfOO#kDH?Yx$%$nDwL7u`$_?ID=|;*w!hzrg|EDx{&2 z{ILFl z3KRzPjIw^#2CrdR+zBGL6o3P&_b=a`Sp{XpoE->g&ZH(L#k2pIOO?)pd-TZT{NT;a zvp;Yq+B@!t&o^%+Bm@NnNCcKn-uTWN8ku>H!J>(pBd#`k{V7*?z9<6B0ts;3kP}jZ z2tHM11LF&zSzfn!zL1Q~t+i}6ja=omHAGu5h$OM4rRBdHm7Pt(z#vP7(sz31l^&gS zC-2=C*9AQR$QJ2f+~1V>=?A`4e68?78WI2@p4KaPyPVbf7(_E|kn)SG@2Q9z&qfEl zEr_KIw!Z@sGE>w4ah(9O%xu^@h#u@K8dlV0)|Yqriba!8D<+-8P}S*KtcfAfT)m0B z$46H}02m+MBPzYitgNkBN@u-P=r}b!XVzMIIH{aTrI7d#V!t&!J+1y-r#p{>QtSChDt*V55gax|QDD^n7~a@C)iWNexi zjoe>u$G&tc`m}-j@NVCK{|;1s>n$=qUkZ;)Sqr~>E>&8JYc7Y^T!TpyjSmCkum(1J zW-g0dc_QdwFrklBA>(bv=he#$2%{V6o;|pfQ%i2_uzK)+Oc&xW)8%vvyLmw@x-CIM zxYAMWnK*4-^7(W)5C{%$os%v@0;ul75BQI%bh2FCH4{|fyMSy)=A?&4o*1-7l8mve zbfPCkZ+oM$DyE$&wl6F!?EiEr{JrY4y*a?jmXMLAy*io?eRTWdB;W2)>2dhp+XOwm z_arY4eT>d(y<9La0TIo#%4)1QYfhkyRmmj8(x*5SGqiuzu2xjJG%o#yyuWY(S1GZ3 z**YpJYWwHU|6$VS8{Gu2Z@@6%DnGw#W&=RFd1Uu!pgi32NxA{e*vsILZJ;S+$yJvYn|+Rc@A-4c+yH6;45> zwt^oG$KBs@CRXZ5Pt^^`*TIg>&jr#fuWi%JEuAl8BW?@zCEVV^r{^#x$H;je$cIX0 zoM*H;oO73~bz)2IAoT}Ad=a95fSA}+GntuOo;eVT{uU`j8dU|>%oo|C)9eI^V zTF9Q{8|~-ihJl%>2j`ML);3kDxLaOV0+S$@<*;_^oKqlY<6+9nzt16|w7c%_pI7iD zs7@-VKt@i+(vo66v$x8>V+Q-!ShTe@HI#CCE&;9{dvE_HA;#bB3TtoIsyFpPU)1!$ zMutB>vR8L8F=F*FUr*nOKXwO}CC)947mihpXTy@Ia%*Mj#?e-DqJ$9^9jq;)IrD;R zA)Q)ZI={5pl?l5foq*r@1{$KWf!x65_8KO7+!vM6G?xT2PRE|zEdx$|OiuiJSu#@- z=KNI0?*{T3i{6gRlOG%B>GH6IA~LPYn&2?wXPJg256j~>Xmb+8UKE|&G%VamtnHP~ z?T@Q;7Nq}Zc0s*bHC!Fr5#DlppzMcuioXaZesVdE8ble<`7A=F@Q{%5qN8u8n75ro zlyYrrX8V%a|81Q+FZ_*Q0&cn!GBUH?O7zpGx92LiW3{y#)#iyT#?Rw>vW7S11Fe_8 z>|rm9rg7uh@5`SbUk%}>-66lasXx5GCXu>O>RTd=Rg=IDi)(^u)Xxi;^8ldXt9+$w>T6RkEOg?hF(K~_T$ zhjOpsneR{K9-~M8M}AFr;cqF!O8g4;0*fruGv})*y47X z=mc@lt*W+iP1%mJ;#8lK0lQw0)eq%a&714Eor~gCboPnHrpm3p;yx=Y%fhN5JpkJh-#nyXP-zX;wZiClN#Ed5VoP;6XtJ5zuD2Yzu5g!#;-ypUn3A3l^ls&DjWf7)jNB_ zZJd6+edQTJFD06m2dFuN%aQu08sC4vR`c*jw%=Sc%qS%v9va#mj1PP>3;Lx8FxOMgrCZ@WtE_SW8Y;w@C zwGDngTM`qPyNr>MAS9vd-KIP!XZUvMKfTbjWA@0S%PpDfO~I9_<4mfnRj;d8JKHnh z?kxTJuDL`Ij!PP2jmx58W`~kDz89*+l&7;<r929~e-5R| z_-Z&J^~F+vZYGj&DHl%pwy~z9x~93h=Ir6QUo?|iuB0OuC#TnqE7Jxb*#(1mm6&F# zb><)|(_y&mPWw$>i{E@lM>9u9MFWsxbMm4o2RfY-6ECQ!gL{7}Nv&%+o?bB< zbtLC{NeBG7{TgZZJMOcYw%N}gZH;NENDiWOA4d9uJK#h^gn`1A1$}3$i zx>v_m*O~_8a31Xs=r{99;bG~QbvE$DUIg`S^=>3~^n~Ta;&Q7y?sR6gUIioVx>L78Y#2re}rcc?{sO@y>BY=Aq}6koi9rAoI=55(EJ) zsT_r35=wF_W0aVAU#}}CXDAdpSEm2eCFo0kmx-tpRbutjR9aY=CLcaPJO2KCTexfd zetiW9Rtxh20%{>PhPt{lTlYFyUE)(cz!cIiA3ikE(P`1p(t=#Z2{iQN1mOotzE~!G zWD<6wI(A~>c5;Fk=)7+CM-gO_&rtlf24BNN7*D=z?T3Su-nNdSN!PLj;gpY5eI=N$ z_zM2P19%S)W)BYuK|u*H9SWqVETklhLPN|$L+cZgogiPW%gVOj8*NCp&o>45`ntu_ zf!Q4$7Z*}se%kUk@B4!OG+6AGhv$bGFUi~c(o&yeP{1QdB0W=oiPONw;k>fkrXeU; z{n6hN8Gh#Fi^w|LwU$`w@E<=?BD%1mW9)VNu$c13OL(gddTNU`x-@62b}#!=^Cj2xxuC??n9ptC z^tiM*);Cf~gLquIROn zFxJqQw!J)jw%I1Aykb%oeC>N4IL`W(yK{wD(fTD|*pfD&2l%>&JUtnijrPJr8tOc_Yu|U> ze8EE4rC6<@9q1$&C=`HT{8xtSovgk#HwQ)zl^FD(8hmK4-n}*7e)ddT-%xu7qGu)a zqx-{$gQqV$g(QcfdC>6ewOA*B* zVDo#ne3A2R4H!5*OXqrbFR|<9D*PCV`*Vf4d&W&@e_sEZm-p&)XR6dlW?^p*bRBor zIPGq~v#_u-%>7hbOJry`-rkP(__3jUeuiohJ&L&BGtG>Mz76In^T)w$1W($CUz)IR z#lH1_n_<5>DHprIZGV2#d~4Bb)l8Xr-Q(qTwOfsZ3kw&x2o=b9;g=W~Seuym=pr{R z4h*{H1!=yuMTWsXi|qJ#s;tJwWTx>6r|Ag7k?R|!L`HMf%^!h5GQ1|kWt0k;Mgw8U z&*Vu7Da6ol8VrXrbBnYZoHbM@!wwE`SJS(jBa-PH*1B7VhNL*bypz@gkG7 zluEM!AX_doRrKrIVYpsBgcE#zmi3ydwV6_}o7-z*C?EAD~PD)wt-#eMywFpPFLhalNSnv#c>jgJVd}P_|$xx4ZG<$KW8zp z-=EKlA?!%Qcd>UjXVEgZDDLg0iPake^MG&~n}>|`VGkE`B}d{9|W_WTL1?A`+YB1?r`M@7yL~=JXPpFBr z+E*rei^^+jK?X9dzu#kS?uLiQ-_<1&KZxgP1l&*z>?RKp<2j4_xlbhVt333hgG4UY z_j|LvFJ9a!;?~+_816^6vjUh@YvT>LOz>#HRR#hVuRR>hiw?HNRB%VVtD;bVn{$BfiSKOU|((FFH!P zOJASqB?^*=@NVd*$wvDvP=}s)2uD7Uj*pK~qr2YQiOEd#*I(bsWROx>OuBvq9(RN! zvf5K5nh5$MK%uzHE6kHKZm&zVJ+!o5+1PB&6em+P_4k9=0gMTq+u5m`nR0+s0;m_E zo|%OOxL)+%z1~A5&L0($DW&w+g?P(~iRe(oJbNe~DMp1BJ){sTt**vwZ1b`&ehBCtqB>ZO;cP;qpOR<&v6Z6kP|LBJq-#R!h$d{$Ob>+ zr24VEP6+3l(DBlyR_cXqEg&2ip)o%q*Q5(>l@o%Zvn9WSB3+94yj7=2g^8;* z*Y0)y^XQ?Oy}dmJ>R4$+#_(3bASMD*Kvu&2?lr*4oWX?R{aS}aUYL=oP1iOC3rtL6 z9GPHeW?&%aPl7OFO6ciYUN|3J2(S#u?J_$&dSn8DG~d9u|BtZ3 z{SsDcO3Flbt2A%b2J7C0=?0#=&fq}fI8xfx>kBk_6r*Osg9_1^Q$IajH$53)1A*97Rqg4TOtpO% zi0xlruOcI3!yoDTk-Zhy8Q1MP(EIJ1qRxUUl|@KXllRF>-h4%ZNbA^gXMSOV#e)k3 zA6dUH#lAHQLY8MA0TVUenjj4&7Z(=*fZ1;j%EutS8xdLx%lM`7(pNUFHps#(9^TQE7Jq5zTTknXw zvn7sH4s;^AOOD-maP*DuIs$f;7$<>Xdzez_aL!Bes1?Vx3X{Wf!$@B>G>jU+PuoU- z{B)|uerx^UU{tBF^FhYD5wf=7_{azp;=d73*zM zeNvps+U#U`fTnmKt5r8BP4+M`vG99>`d3huyhRpeb@p_XMblx|oLqAJnF-u0RO7uL zR-$&R1iMs6K0a`8$c}96R1j7OZ>s7VgNMUtcx492?G5SzKOTYSYnG zE6B4nTZ=xY?Ey|T=}L!!Nre)#QB!BH0^af*&nR} zzG||*JFel2M+(vD&4#dRBwQ{4k=D~&NlC$h_B{A!weIl0a@8sd3enNgE3?JPkhJ-B z;owz3Q(`dyBHU%efm}I?5q|_qeoq_&sw1;?3&+z*J`Vss1VkLBrmnlwtpDS2Dk@Ik zFLZQt_`=y;7Qb!JYWD5Vz7R@2V&irLvo}$-=H92J8FYjJ&V0|7-x-uVdQ_B$U`P=y z4b9#{bE&7xKTG^z_2wT2j7GOJPzh?)g}8EZ-PO1+D!F}(73PhMtdJ3NHn)*%oap;H zUnd4#`{Z}^&7oj*lRWW${cXv*&VK9Yt_t@5efI|)0T-8x;T(y?hwlwc>#QzpYuhLG zo@7ep_jhqNn7$GgHTs;aA+CTBX~Dw7u+^XRCD{RI$_26LIG73G=?&IE2KR{3 zq(VC^%!J!1DKS|?O--xTR`{mipU=3raOb6dD_tHdR7qyDU}9k*QKf;`TagMz>tb#4 z5q-WZrsW9(25YCUx2UfVjGUoO-_V1JQGRj+<`KUQGNP|xn{{19^_*Txwf5)nyC48=L=6N&M6_P zH9HQFc)(v@RaMuJAc;CTv4{TrV>#UMf4`yHY9Ti}8v_FabXd^^tPJHyz%HSnu4t?S zWN{CdW3W&H9irMmNJ3(KVr=`39*;S<>ho*6_dpgR$d5r0z7-gs*EwO^BOCA4r_)%? zgj`+w`};wabbI2DQ90+kOZ0jYIyg8?lxQ6vAFo(*Jbd%y;yT|KZpE zMk8U;gqBYja*D4X3ELVAuk(&}X;e3HI?*+Fd4Xn{?Ck8YFeI4QpLF$Sv*&km zTRS>FUZPPgt)ikr@cabz$%(xKItj_Bp^+0z=U7Z^j-B%8=xsP2qN7v%mwZ}0WvWABnKX^Z_!* zp9(s?5V6@Xu(6Xdu(@sUV|NTpYK@k@j|Zc3-}1YXWLqX_r*g4UP>7MB_Tw@V5wa-K zJTH=z{DEB@>DfTiXuUL?E7f0B#p`y0vAA-4aSI6gm9&x+y=L^qK zp4?d*``=tfkCE|;igVVM*4KZo5DyJnf357(s;^>XuRgxb7ZunJ%irFaVtnxU%hQL< z{s`hvp312xXz+uJzEFjGLr7q|DgWa7qPK#@%4lQ3_2kLVk&UG$kKthS@SI$%pFgp& zFpJg7b1N%j|NH)sZ22uvt}ZSvrtR&6L&RF$-7PLG{JFc+Ku(%NMlIp3l0oEfD@5(G zeX!({#@pfdv5hz(=%khOB9{79p7gFx&0tseb~PR^QDcMU&g76g%zdxUAAhAHOFgBmjE2r_n*2P31Lj?O}Vza%yO!l7fx zM_*rGhlGTjR8)!v_~EFPt+vcJhHmsbz|yb8#1`yKL7L*c%%J)Dr;eeB$W2NL3oDD` zK=K49Hp;M58rRVgmaQHB?u-+_GCXBuBy)8|flRjd((tjc5OZGGt54&%`m14K=8Kcy z;o$-H1zKuqL6U#2%td;?Uz-FNQg4=80{~1<8VqtLBfEe9{z$GAHkbXE+S*$Gy1M)( zPvFXqXJ^;Doi=cIHV~5M-5ePR=iSbQ`<*nBs(;y-D7{_NxkV`~syjfQB?Khk`r3`t;{at$e|9?==8gLPSJFa)3t-3}gg9Hn!8VGkg}~h~VJh zxT@3A=9}F1oBgXUFfml-EmyVW+{y#Eg8r2Jp8h!PnnBNtvYqZ%Z*A!uR@WS)gJekP zY!5>7?O6VL8mPQ7c5(aJ zcG}L-T9rh6(m;A?S9bwO#Z5Q!17Fkw9xXcrlbN{xFJN##+!GgiqsT$3D}sAYU-vjux}S`FYp~8lVkIq-#@+BkV!K)HT~}YZ&K{Ro$Wjz|0pmt zJS;CO8xa+Sf{LmSuv+^1Li~sb2z~%d^cRo$HBwNBYB1WDHrXE9;fD7j%>pVsix)Sv3k;WsgAH<@|VyFz;7ldMC3tB`xZz|OIn&ZC|JxkH`i5E2$HSvkNk51Ezp-VkD-frhiF?r}5rGd+2ZRH@sj2w_Aj$Sld4pn3x#8 zvN76JRb;Q~zt51g93OvEQqqDi4x(FNj{uH`m-rGIjPT@KDCpS}3>=&kjF-$yb|#>B zJsr^NZk#b*nv`}ljixZ8XE&0fu%ADx{Ck28)1SJEXL31w=jSHoJmS>KG|kv;W!Up zYE_X=O{#Y=(|W%&lay3#%Xv$P^pDl|qI?Lvun#gBvMlWE{2rGqj2j+A_Bb7OD-H*T3Vp;5Y?jgt*LS_kohr9j@Rw!3_Y)_JJ(Hj<4u{wL!UH( z_Y(@pbT?$ruc=Ty%`yw}Q}@Z=45Wbpu-9}Z){y8Da)-A>k)~0y@(L@Yhbkno%K2nN zXX&5n>QqS~&G*0xEWOXk|_sJb%Pv$8d9I%YNlcv3*WW}-yPwAkaHI76_S8!b;u4{XRh+M7n z3r8le#T673R8->AOirTlLfZES1qXcez^^9*V6`xCIe>#(g5!i7?%C7?!*QJy5fS>~ z_Yn~Z5tsGdx7ov$2{6G@z3$|?$^uzB&D{=KaI(}Ic63x78fs-wR_$C$|M!V}TR;Za z&AoHfL`i@(bhtX*ot>S{5(QIt!Ko~_l`DZJ2XY|(yu3G^k;LGp=IiL{?#^6bW@g@6 z1o6J({=5{re40|BGJprukdl6+O99~B7)k{)$%NE54WgD)wZI{whUqe zsZ^4(JU!h*Lh4>!S{$r)vNAYwwB6oVUf$U52FRsU1u(ka8kjBaoagkV)>iwU!n4Kj z(~bF?2XgPnT;DX4ww+%gBRs)4$)RMhb{6Zbx&MG7G&BsxOJyo18xa<#jBHu46+a2P z*;w_61j}enDGgu53zde!-{$BK!|0sIxP{28Q7=B@rWNMMcbQ zc71SlKfk|Mb7NKmsoFPrm+IlvLMT#wK4x7T@1nDw7ot*d`T{y9WV)8lbCQJ zSAmkn?3^5n`G!P150Gc#wqE@99*Gx4CY8$}Nf3GLcwe*875X^u>({T@H8tyNJ(uY8 zN!>j&I?c+XqYcG{SIn$zO%@hy%6~33FN|N8+9Qx8UbU^=@v!#`qvop+ecQ6-EZQ%p zJ=t$!VPUbv0tA$ejg5t6bZiXR=_W914xreDg@qqIdcc2bJ12 z^U1G-w6qp$-Laq{1UNEa)MGlcQ9nXI5ZiVe5byfifjFSohv>zN7hrZKVEXd(^aLaA zfA$<69ogF1>CLohYHChCF4JywodThXH6ZKl<1h~>fpG&Pm9+nLw+jEqX~hjDDM}!T zE+o$XvBY#z(&heNh5nkF-LsvY?X$D(y&g zpJr#{+O4~@3{7QcYuRn=rLpg8$hSARz@gJ{R2UDT=d}eR)2rD!A1(vWL!X4aG>K>G zd*_#qbhoHSxtEutFkVVZx2Wj)gfD(t$LqT(+-DC==6Sd(gPBOk_a&w0ncPmNO-A_Q zlw}gj-t{KLAm`L(iIR(mdNG8C2f5M_v+ z-T%7r%EZh}51BFZl1EXV?{YFaVu}CHWY5Jb|2uEOUXO+&17Wc0m+=A}6zl7TXxCqg zi#?9lrGR$^vVC^a;^JL;{lM(*fi9&0JMcQQ6yWLuQ*7Q32sM!K`P>f{TOPn6?EvOE zY@_4#KJyth;ETx1qs#=efiK8c$b-T9H;1y+MguTGs6f#F&HF)CW+u`6Q+RlIK(M_0 z@bTlz^K*(7@gLhTQv2k|Z*_M7$p;>m8-gRG>q|Cqxc2r;q-7DSC#k6RFG1!Dq z?y;jO{cMyam=M$FH>^yTs!Xe=5}q8GLkS-D`cS-dLqbhFU6?KVE$S;E5J-8o3fUFV zqxu?-+PVTV8G8D$(TNUn0tVT{zkGcJTEBl^$Y>Kci=~}z!5M=}F0ZUi$>Vc7?wlX4 z81*K!SU#*HMx)d0w2X+QR>{rD0j-YZ*hXOgz#s1OesDplqDBwmIzhLU93N;md!|~7 zIH~~9aP^ljpXQ2J!Ex@Ok-d8L%4$Bpq(s@DbE`exc9Zing7xt6eTd{gUdy@ZU5?=1 zbtvkykCJ41_tSvGQBY8zvflu72}JF{(u4yK6O^L?5>$;R@XpGCE|yxVva&J|1v)CW zMG}iT?#P5J=uilbsTm4 zOa%ddZwhBL|9A>O+0XAVP%9TLN|8=gn2dn_6D6Y*AVoh1GXx9ZSAbF!P?N{I#hZJ2 zHrYBv`uO~%?jXDm(xUP4@zg4YKt4fgk%iv2j>)YuqEck#t1utnViE;pAB7wm4pmxv~piP1o5Rs=D%ni^)Tf~xA~ z{yzQlCqHy;()m09O0T_i%T87Qu-sVx+&9O^UlB?Q;km9|8Z0`@janMG6g`Ol<44|e zr~S2p#STR^HBRS)ciyiFBE`TRrQPJ-HSmu+55ghwUnlo1JRc-C2?@P-@2G`xy^ieu ze9Ap{l)EeD-hp00etsgK2N!&j@9)`zj`tIH5tSn}aOZzc} zKhM0U<$Y(+fJI|UEXvGW1|bqACZW@OJ=w*okGU?f?Xn1;v3q@<;lYd86Uz9>=^ zmzM5>5BxuS_4?P!>PkWT06M)kjb;r44NdRpC<_@G;QAMFsRBtbYU&;Upax#w(HG4# zWx9A3B_*Jy28M=v^NkwNyG8u{BmK2kpX;N#^g;B#J(3qLOM&|2A&3n?T3I>G!pf@I z`7jgsmEPVh+hc{X9j~4VJ2*gr!hkz|XJn>oh5* z0a^I_CoAy#a~aVvcgG6jWhq97hQPdD&^Kz32eOXE8r8tgOG-*IF)_u&gX}5}m;I+U zovl&uHt>rz`d6A`ev8;%F$j@ZS(QXZMTsK){QP2KV!i^Nfa?}txcT{ckdp~O!ja-DHMZ6^b{TgTdUy3)l3sQg}fe!$=GJ=A7lV9JT3{-2* zo`BC!OG^WB9Wdp3oXN?_XFF3_1LD%smL?`Hr3C*j_wQ#x|32jJ#T*bv^sJrjFE}|mvN1Ey z03LleH!!cJ9(a&|6dU}XAU}Q;clFLxO-)T-KO1<4fWSgS1H1X(-M#nsdXO^<3QT8f zY=Dk}aCU8NZFhG!o=&T*ynJM0LZ{B2(!kingcEeR#l^+>`U(OE_15*zdOn;HCOiZG z#s3Lt#9b5V_V)DDR4T9Aoo;Y*&wT&x69yu)Jdm->&ejJPk1m9Rf#Kldc;out`pUie zySyg2w+Qr^rojCOL>BlUFJHdI#=>G^W|oqY+DZpT8E_cDA>lAEGQKwg@^d(c zuh{V3`5OwVE3^zP{gIe?71JDX#1Dd7tNboX2sT=Np)& zD?^hv{=)f9n6py&G_Uzr`l4|DO#@LZUjWAXULVK%rK7VmNQ&i^uyJ-rdwV3`MCN~w z@%+wzvvojIQxh7Z%Xq#hMl^vpB;)xWre|h~`>uM~$p1|TemwuC10N3LToikB7d3kc zhiL*uL`0^3{aT>^?>F$;@b9xrM1Gf*lk2Bcb8&WNQ;yZs)O=U}SL3Ca*!7R9|9bhM6=`+&TXIj&n!;UaySqxLYHO@{=HJ@A5 zN!R1y<>ZuCR<5t9v9h-Q8h4WisUBqu*I=-+vhF?m=gIvCuD$fddeV<`epA9D8uLP|@WJY@w|!Ezp)u;*X$U6Yq3A(v|&| z>ndOHLS$uSbIj_Beb-hC_ROq*GN~XdyDLx%&jG>#zRdT*dCrD|sOuWfTjz8I{shSEqPz=dH?qd^MlHxY6(VZ&Xxnh~UH51U~@Vis$OG zyQSqY3Q#*fi*e;`6sxhWqEH`gu+?C%dQ08kR20_Y6xYWHr)Fl_TUix@Vz65YZ?I3g zz(F7os)F~i@7#Rl6)}p6F|C5|_Ixtf`}f`5`f6%UC}?ZT57i7o96)=f+>gw2@nYY| zh+MEXEFh@aFlQ*91!q4+zp>qN8aaTKW(#+|`>*el${82RzeiS~(K)`{JaHO3+IV^E zo0=X_s-d-;|1Ur--uRotQ_8Lj(#^LS?yc~(uox=5{I%*X)4u)tpQ1>hG#VQxWM=Y$ zk)R38$#FV&ZsXb_k2P!~)zyt((0jDCwoW=N^8+Sg6cv-?+Q-5-3W|%zu_XZbU`raP zx|BFjtIM_J`!g^?uE?>css858jXrlP9%1dg+t1j7f|Ji!>6F!Jf-23-%s`xA_VyH- z5-_`#OslLam*^BsgIs?@SwcsoHHiU<>afh+WvcbX;#(Ncc(9=!W^V}&4E*w6EIhtpfV``gQe_ z&G)J=owV!9dhx>e`snxX=hkejtxFufTtn4C!39fz=HB$H;evnID`NV4>^)ekekh;Z z^d?}3+E~j*RE64Vzy;JP8jjz6mo zqUTlh1p9>YwK}$bH7C9r)is%6kHf2$jt;H=d;Z9G>L6_);6SA2APBRwkMr@-Z1qQY zAcW%M(w zE3{+OR9v42_+Hcm6x<<~SZ;`R`}3JMVwc|E^Ia8q->na$7v!#vUW10cN>TiQ|7{5V zbv5|?)b#dV3Ea)D@l-3@v^D~Ln9-Y?-RRo!T%gQkym`FUUy?0U;N;2n zwl-C6oLOQa#B#LSU~eiM2WsG(C~_L)F(|z<_%Q}h@SEYulKle1qW}7za&3l&D@krS5P^1`?3pwM$%Kk&pqR6$hmV z=73i#3${aQB#@<%>ahf964(f-mWr%Sq8`;R&rL%hZ9gD#p33hdl#8r9BYS)0XN$!r zY1FFmq@F+LVP_Axoxq6gIolmHp6N)_Vs~nQ1|#air!*zamJc}lSnq-;{l+={AfZvRkA zU%q^tHqY&+L6j5}i+|?CjL=ykEYnp!#J@Uu;)J}S;))cvVYdvcjt){b_^FhQPNVcq?F}rbNgB8IPu_1P4qK{o?MuMnhWb_3`5t0M$@UbBT%;-d(rMfc{wtqlEL)q`=)_oSID@wo` zSrI2I-?p>dqMn87PF{XK;Vc_IwY$4JSQwlXfEycLwm1DG^jGAf85@f@tAIF(W7@+z zr3^0|d~``k-;mKk1`&99Bju9Q8PbHTTs1V>@Uek9dL1+Ey=O*f2U6Feo6yzJm`%Q4 zr2J$gEf+@x_6RjvJdKR>bmJ24l3O=1q_W#)zq5-=s$z6oOUqj&mD4SsAXDk_a=KETPX1oDyYw-DQG@@JOiqw5+5MkzmH z)o54O<8cX}yKXLxrOkD<AR`P%mY+va|>^MX%fWi>@px3I8) zE(T|>^w#zuqUe>0gQB7_d;$Lc)O(LzY*OLui z?%**H)Idl%8$s$r6S?vW=-q|`J{V8|#RqBzmKLAIGq6v~ai1DnNLhGl_PBm**8d+p}bs6rX1VYpab0uTreo(1b zcAZ*er(P$(B_IG_4mTAY*TqJU*6?#OUS3`ncP+|6MP0@ncvgOV7INM03scDj1#_FO zleXR+*`-eOL~ciMRgM*r4Mt^H!K=Yxm*>1HH2-SuM6*Qed)DwXr-4g7z{Br5Y1g1S zXxJdIm1jzkl?%Q6xA$N|@1rQ@u4aNMT*|p#pqJdD$mm+E4%$bqF2x9Ahy2|Kb0fq$ zY*>YACEIhC({@HN>9j(`v2mf?xG@5_o+rLid3Q{BxO4`OkWf@g_;+J`?+|A3)ei2B ze~y+oI~j63@kF>(f@PLwQ2{#%FX}?-X%uY~lMnCTpYob_aCa|~j8B7$1?;1i)E%2I3PgBBIGu%XJhm9{{3A7 zvUr!ugsbZD+x8&)0F0|aA4%x?M(U=P7Vv4KbxJebXy$7M8BvDXx#JlG_J(uv0 z;WV~9llB!zzC|Rx7CJ2<`6IrsYtQX;sYtb7AkAzoEvtpz-;7y34S51bZIv?a?BO$^ z(m7QLMa1_HSmoQm_+ET<4g6I@!v#pd`H~##6&V!wQ}YUhr-VmrQrxdEpD23i{Z50l zkYa_Pdy6>RruO#YuA^E684!a$s?yF|mdpsx0%~gHtj-M@G;e6U0g2-REp!s%$HHq% zU;WUM!@b}5*=H?5ys^f8>RWXLw=Oc6#ig`KpoZLKi~9li{ZfWy9^nPy5#PrFbcctB z(Lxb-aeJmmX;5}kwVW?@9aOmq(49$q*m@>(x=S%Z9}aq!oT|Jdt=-#Ya$e&aWho zAd!Cm7*pH1rl{Gt7d z1KFqORah!$p%W)Uy*(WfRDhm9>o|C3SN9jU)z9S)ld$=4lI;5sFEo4A~IH=Sp6C4~zP8V_oEoU@!YnNPqn~op zqX&mZajrBn4bTzGdGjAX9$$4Klp_ZlU)5#2iPU;3v@R7n038@Ds;Q}p| z(R38p8(LU!45Z$#09>|u6r(bRyCvsW=6Xq&X441ZJ1Op}skyDABu-8S$^TSf)ej>> zoA&{_F*E%qmU*~u(#Z)wXK&u4N=4FPY!O>GZu%a&v1V|I+6japtVc)t&b7XJb+x$E ztS$;@$qpF+&=5*c!c@%T$7%np*($R0r-JTj(Tiiu5dY=#UY_^7PIp z+J4J0{sXw-&Wfk5v#+frE3(|<8RBV%wyPRonCXm@jJc?Yh{XZYhdk@{7c?8Xwtw#E z=ukeHcY<>6+7BQRRddqLS(1s)0se7BZ5^I+9Rjgs7-0 zf(JC>uui%MMAP(z|8q1!LxZ>ls-OwHgagX#xr4oP3xO^OWAYI^`caXQ^t7~<*G}3v z9bz)fwb<1;PU;i%lS^)b6wu#qY;SL0^L=n|ZZwp6Lvln{g|9dp+f6=jM@_t=t>rDH zaA&5I_w7Nvs`-GbZNmEi1ybPHv4-a6eQ3Gy-iIIPY#%V{yo8wobT3N!9Gsl{PCFi% zA?oUS4mHG+D`|TA=fca0j~+2}GFNXWL`y~f{FL(!onDkw9x#0$035j?;zPD4$%4{J ze|4-zIT1q#Z2ehGp@4$?r^sNiAnop5vSWa_UXyqc1rmZQio;XJ{G^SBMUJN6WNVIr zEa0UURt!bciB@?<>3ZS^Wv6!rZ*?GAxr`A1p;t(Lq-r*FCq3~qntX}%lImcaY(D1( z-UV-k1pPSuiJ)4h_zGLcM;tOOe$5FU54nNn{V39&SFVs#=cy2e6KqG1S^|dymn3ZD zOsqwH3B*O!85TOIrs+Cd&z^%NX?EURsQ>WrxjQIM@%h;XYS9=m8n{+Dy0?Az^0r9XDjBsjVBQEE;_tzYO%XW~)u#z?fL_5`|*rmBfV$ z4Jd<_?nqV9QHuTu-e!V#ZTDd`#9>!QjN>mb`WDN8O}2EjtZqD*#nW7mzO11%hhU*j zwsJs6&DGGhMl~f-O@}%*iKv!2KQ4Z{{B`5aF+z&zSyM3OD%{SsYS|E{Mt>f%s7(L< zr(N=k+fOF%8mO~QzC{)r85&ZE6|DJg`IW?oQVkR`aLK@gQfm?0xB z3j|!;#v8Z!jvdn~bTAq``AI`vy$(Whl&zhe52)=PZFPuzsE+)pTv-L^C2^|+-QDB} zH8s*#VkJGL8f}+>dA@9zj?#MRGMtgor`kE_fUl52p)WMw;Lc5l=KyDOFLtn@++4VGDz!I3q3qY(T&lp%J$s0%s#AzeDBD?c z126WnXOKLkoT{xff&{$p8swb_rmlSge%a7(}UAds9@%ojV`u>r-oKyQ~Y)iYvd7)CWVj zpgSI${3&(j4V?)0c^ygT0K%!TIIO!K9UCK6fbYt^pe+xE8=E>!ZEO_saNP3>%<33* zc38JRkAa0FyndYCN~A8us`I&FMNg_Svuni>c6K#delh2Y21XP6dnp?^fBof{|LslB z=<4di=7jR4lqEX*OIXG^8RQ-bAC28k7R_yYl1D`Yx zW`ieH7ZzZVzyT3;8{aPwXzP+35AH}oPENIo58~61&aMLW$gQdORqXI$Y`XGo_YIgo zgU5%sg6x*6JT^KCl^{harbg@wsjvLUX_PKHStgMy6DWIHW+NuDvRjyp$~JG_e33U@ zUD`j#g0t>uUDU`?!S6OtItIDX4^wvKOS7}%Y6(|8*uiSTIq@3MpqVXEp#xdS-nq0X zTk(wzJBy-dcpFfnyOr22LIPXkrR?*mQ)LZX$2sHqnQZJif?HvlAXeDO4Md`AkdbND z#%j@?g461npLOlLB+s7J{$-(l4uIY;|1SQ62!9Rg ze!i{attDjK7ky27{8O8jm64ZM>Ag5VZ(R!Rk_;J z)03OKrz@X0yJ@F0H@DL@6TqWfM>GsU=ynSQ zN=dT*7#;@xcoZG2c9pB<{3gI9GzYZ?hd+pT&fW=Ecm%{u+MPExtr>0*P0mj3@adDHdkPrmKh=&jNGBQRzdL%1DmM(Gb37^{9 zTp^)=uOmk7dWYt$x7W8WN}$w%G?n)pwFAKZf@pHmwiMM(Xzv90`LGj=a-9EgQ& z1Ys`#>R4@z@6nxN*RNG@bi35-+`nH0beEf(8|ZuZeTpAfSwZLF*Q?pfJ~M>M1j8tl zmLo^LeJaU((=-Ws4{cD|T3KfVs+J=<)#L}7wTneWMfKf5Hrfq0CG$VzcOwcbcBQLl zXiY~6ow8^+xM|a-S$9p4vM@OS)A;n)ePDj0u_rhma22GtyX_J`oqoJ#95^-WgrfsHv~DVl<= zBK5ww$25m}OHYr_m)Dd!jzsvy=sPXz3_NRHO`50XFJ5`bgQ~)|G;aqKnnrIE#Lm5O zK>zEHAJK7f^;6AyS?;iO5?2&{rZ#?fvCh&)ef=!6lsVeajt4(F3w~}oxk>GSmEa69 z*wcy6({ci`zy`h2h;gO?b;t40DjfFR4o1X$A(mP?1$ktL2M?bL&E(x?XmsXbF=X?MV++>FmIalMo6!E#b$JN@@)McpQP>=RG z1_fG&Nut#o0B!JU9gs}>@z7#ox^aqe{L8(UUW~<~M}}Yy&ND1P=M(@xj!|i8Da=Fo z6zD>hK^+rIteoX|B9YeHk3ewY<|1Q7*92@TfT(PM7>S#q+pw0}SVr%A!Kytv|O zF^3YXRzjPSzof)vOH_YbmNSGUhd5M`PH>n=L3`>^)800zhHmU@?o?t@F+7R|>e*iw-MssK%kQUVH~ zfx5c%+&UMHmzCBhZ8~1^#3w+)KF_qr305Qsw^|2WYtbB}q^7E6qzM(@{(jBQRp&e|A8j_XU^-|A-xw%9dxbuGj9g7XQ(eOBZEBLhP)ldvF?`SIhIk;4J! z$a0~wvU%pD7nsX5mBbBQDg5H%;;j=Q3T(Btwa5r2DnbJRl8?Zb(fPwc10oX!3Uus9 zw4gM?;?;YaWpNIVstCn~99@Yq6R(gDD z5+v^EzUKgnv$M0f<2|)lt1Cz@u?Gm?C#15cYH-0&idjAhJ0$DFv|nfm7G{F>5t^Cd z=74iEP|yHc$*|19;mK!!zjJr?UY_iio(UNz!=p(ZHKA0CPhV;GJC<9Z<6shtzYT%;?*7w~CQ97# z*K;Y*A4GC#OH!MCLO%aA!2H969*49pMmW-zGDYCmt9D*rNgzky>=~7}-c>HZ{OGDp z><^_Je)?gD&!w=>;?`|kyXLq%LA1C!SS+YmN{$D=<2QDedgU({oJ6e<~gpk z67&=}E)YA;)=JNo)BQa6Wg4B%g#Lb~;@c5Ipq|vcy}ZDpcwny@o|_8Qu(^+q*t6q0 z$65%#b2-kJwju1$=jFr38Lb3^DPj_4j^Rq@*%(c9mvmP(#ubwDBJbUESX=cj7D)R5 zhk)hf-W@5)@)RsaR*j323mKNfys^LEc`!Kf_7soskCX~F`Z>*l1{>ev=oBJ79CWLmwUVvb7 zPe@4x9(5^#a}5~=ge8rGFT)w`Z(~1_cfsgky8ZwIL$!xKTs~;S2-@$<7=>6^t}UK2 zsk#G=7Py;kv9-NDb-5<dF$FQQ=XdAp04W>8O+Ug)+ML zzunTNCb$~*s9V5UrnxCO`WnHQN*Mcv=)}8akQ+wFPq|MeJzi=ebY4}Iq2o1J#+?cy zdvxVTed1~5&nj96jz{u#Xa4mjivPA*_j$s?1GR-_-@e*wPBj#SPB~Z=ZA?uEzJ2@l z^QSl#QvtdG%6#Xhj}wOOtfkcfdIg$L4QCDzDdHa8KMW8MY0^`HS^@SuO3uM@nue7$ ziEEM|YzPYi^NVb0Nx%-VE)f@Z~+WhaIiV}tMN@0B9XXHwfi@Xqv;V`I| z5fOAFwW#z!Bo=MRNz5Wwo{Umavr)$+CLYeau?UKJ_{R_V!w0{Mg)lNR_d=152>K8n z7S?bB=E0od^O%OIu0G1a@kC5@3t~hs*Fps)Sn*Sq(Hs1QaCNVBRDH6i_zm^f4yMrv z*yHdLk5UVBpC2*7K+H2OD_{H>4FckI|vBRdyqm4kPxC?9rPmB)lmGQ~d^ zZE^m#p{rT8oq|zJLU8?g6+bo$P6-Sar}hNX3{-D%relR{HuZSea6~o7lDhfToX21P zruc79e4i)GQDiGOGqXKtOq^*HH?-mKJAmTHTqY$IolRFEcE)~1z5s?)2|gj4+_VJu zbBUXBa@N`toBBBMIAJ$CM@K&1D{P_FAYOOx+Vx7XDLXhC3I*&AVHKsNdXIA85X0z^ zfq?-uHJb{ifi3v6bt?=*_Z;F47cGs9jzbRu-HLMsPbCaR@GOG;G)(@uhEiOD^DHU( z8bpiMo3th5d1O0e_&(E=uh^Tj2$2TRg|@4J9efv9aw9uCr5V*FoN5SC=)C~mhTZ-t zuza8+zXra5Yao@7rUrS6KM4SaEKsguAT`a$6k946T1aZd&FXJIAB?QEJeQ*Ih!4)z z0Uo^lT@ALK`Iw;r@fR`z3mRVSBpRRaP*p=iUyNR19_NBe>CVIY!j?@(&;inIGZ{-_ z1I#7)IE0x6unA#+0g~U;g_B2>p_~gup9{Brd+Cdn(CgPDandHgdn6_P?fkK&2aV|) z)%lAMYl*bdZ(N3(Sfkup$4X<47S7Gx1fnj;eG!39Rw+(Q-({;s7mPu)d-tBg-SVyU z%Y??Jge}Rl%~_l$8u$?}7ed$oj`KD+S{$%!OUngSk+U<8Ksk3>-KunI{hRcTXA+`v z%E~|cpBtWgRhIkcwxYbe3Gcvxmkg&IF51|P!W&}wiR0vk9Md9`ciV*y(9%k~?G}hO z3w4C{Ai~L+;k3U1gbu1Nk&&87lo^@D^7sW|OBSNZb5v(%w+mN^r>M6F)v0+7ktC@C z?Tn)Mbc%C}7iY&ZGBfK1Ih%#1Cnw|gO{@t0>Gu1~$!jDiyi&361c?vACro*-k@cAI zk)VQsP$0|awzlE51Epmi`ar%Cl^Aw#DThVwGH$N`oX4iHf^LByWXwH)-7Z<08VJ?EgG^$w3GX5h8V?>zW zlaP==w<3xK37h6teRc-BLS_NaApDNL-g`D6r-CseH<=)UDhL6)DY~fh~ZYVCnNJsNR%Aa<} zFK*{=jHz!jd49FALI!`)&Tbss&j~(0IFGc3Ca-Jl8H7iMyl_O&MoUUcN9cMqc`l`yy*8-tK%BKe>7#zRA(FN4VtBm@R6 zQNp~rwpF|Sh()@C=MQx3?QaddaJY~StTDoFFmEliFRNMDK(K(@8bT3%;IC}AWj-!B z542n4j;D(p)X{NTmJ@DIj#|JY40sJD_-lXv#%+7g%*i9*poc&mbwkkw;p!VVNA>5c zFaP>I+B5HAV`v6cNhGXj?U`VPn~!vy{04ghJDz3N-L-K$ywyhJhl0O5t1nRm}t22%7JDPLj@>ZP}cv475~i1eJB9Jv@Zg`2Ri&XQ$sk=;c~{ zk9URiFk8$F4U#&lUR{*{@VGF zv%d#WCj4>ZoD@s-$B$qsF-0@z*LB) z3&jue=k#?c7R!N#AAFA#YT}u{F;4R&-<(gPpjfk8o@RZ^c6OKGl1r7@dCA>uiqWv= zf-`biH?8Rb6V-i>;VY-x!v1)k<_&)Tl6F!`7BN?s%5p zRq*>?OtwU*zUTkOq^8SyP1nzt;Y+*Pt@5H}dhPjOx~hAUTo3ld2TJZ$o&v`E)Wt$H zF!-n}Cb>NNHqkvq(U1iQ1)H?9PyrKMpg$e_^S3#dDGVIkkuWhk{| z&GcO}2^qtJEs6c0&#gaCb3*YDMl}6+kAzl}ab`{qmg$kB9cc^jUhxrMWplIB+grPk z{E@A|&cUx%keBD;=$JqcCb+x1Tjnw_trtJ4?y2i1oW6fRtej?q(M1nk6HKH&SV|5M zG1^9UmV{X1?j}Aiu2H~r%~Oc?3bg$CN_dKx^?Q5{hzCknsDMQW_&TrT`wV>Vh#NJ4PHkv+7 z?=$K@B!6i6?q>20VACOKdxMw=jHo}8iR?((OMp7472!1_BC?y zt4#d5*?|9!j9;nDdtVVyn}bo=(uSxKiC z^YVX42ns`LBQkh6z^WObLnLB9QVz^lalWgcUcXMjDiOd>yoeE;*cFkBFDU>HVXOpd zgp3MY0POp(Y2N%Jy7~S4_m7@(e(69h6drgJ?Af7r=%Eb#lzQTt@v6!hy*DgwX_&lD zIrMn+8-@_U1SNjawuIZJVdBQ2fYj7MU`cPv6#P~kOv-g&Dzi?Ix{FWu#kh{!T}FgLMbd{grD*-$=oR#pg>DJkXiDML5V zA@z%lWnf?HO(o9qj+3tstn+>i&&XWkBoI-5uT(o6J*=y(t~7LX=i)X@4knT>U~b$b zyZnjfNqrA=kQ;Hb_km!8#Hz^hKDR-s)_nFLH@xF8v%dbh`GoCf{>A3bPThI`{zl@N z-czf(yq9qFo>D8WTQ7`D4|aO@db%ZmoN?L}Coww`_1tghh{yif>kb}Y2KPr9bVx~A zsP?my4OuS&gW_y0(&MY|RmcEy3C`@AHB`8TG_Wtg>`>G(QqX~2o{Y}X&fKo4WgF%y z8_rws=tGkQQ$2iES6e&q{NEL0os#5?87Dyqrvns~k%7~*4MQ{#pG;|hl(6H&)xlw} zfK?c<2^=KXuU&&@4TGG}qF0a*Fb9ir4seLh9l||%y%5HT9-Q4omVjB<;a9HrnMFPe z%>mf&0v0Cr|(=X}xePi6CTN5MF9AsE)79W)oS?ZkOLwW2<^w$i5Xay>8-|J6tp z(!+Yk=*AGinLCkZuai+$B|G(?)wPlkY2ndb-nxg^j`wqVtD#-#H=lOX+Q_!5=o~C& z*>Ctva9Vb!ub5E$4w^OJ&AG;n7d3@ao!9t{(yBKnDn{7dw!3N<>ODLyV7Rji4M_Q} zW%`IL`{Qq>&jF!}AVg13kGiq-wg5l)MI(_N&jV-U&RPevI2@c*f^j!(x5e`jw z&)&W7zo!8-y$F<_Ut`DoH+Do)M)h0BmHMm*MK9u^0jlK;?iW1!0Q@Ud0=1t&??{6Z z&3(vkFj>Tv2kKP#ZG^;*AI~l=U1lb1`s2bGx=OA$XA0gA zOo1f?c?I6NuBWG*Uh{?!6f6WSvL3?;GyF<+_y$q@A<{xR?wyF68YVhY=WDW5#1c9g zvf&#^PruDAE2mli{{8%WcTg5ob?Nb#JcGAo$0?-Kin<~ke|%Q0*Al+5aL#itG0@P^ z?BD-;p~hpJD)OMr33EV7sPoz>*|7bdMrLoCK)IZp+?7%Unu~$F_+@HhqE8Jgcic4M zYn`yBiaskR=d1juy>ydP7$(O8|JPQ1S0FpRmeY&nxMFfStgu1(dU7etsIy)i<9mf0 zZdJVJ%}VyU;;7e_Kl=tM^AO{-3k>R}h7OL{#xUE>j?`qjdtH3>)s)SGnh?TZtGZox zi2u=^%liZ)9c$glQ$r3GpLacvR}7`mrH}G@Lkfv^Y8pU&w-ih?0+(q40cZN@(3a-8 z7%|$i4HBokErzo(4e^ug4pN#d`crfJe*Lbw6|k5>Dn-j7X4R4gdJp(1u%V_Yf!kFyskdO-AoeT^xpbt1~U)f1V=L+I0aozMcB-&Ec;DcAe z&}RO;H1;_!Ls$nRejw)QxOH-~LXitLZP>8+EY~3{fUMSk$Zy$%&Iw*;-$h{*xtVz@ z@Csa4R|Cky{_J{Qt@U3%PEJqn>Gh?DR5!$Mn>phpkjTo)N@WGj_A*okk;$JA&KjyV z+Q-E!`_Es0;j@BvMmmn;1^9dM)%mmYpAN%TeTiAA2mb6E%zZ;cL&!PVAtyt`W|uX| zmzaW99PN}CidNOr<62&9CYVX5MO-jbg$?F}-ye^f5{ zMN!ddblb4S0|G%FKEQYBR!+_@SVRD;)R3)xNo#YM|9qe{)UE^Jt#!+a5(@R!+-f+6 zKYZ}&%_A>d!gOS7Z+u8|XeQ+Nc^H<`d_Um3fPkvhT-hIT8uVWO($Cl6`yo%V`=l_U zLwjFSL$fBw4=R+CKzLgruY_)P-mQM%E=En4C9bY9l890F?_b7S1C`pC3u5mj@0vWe ziJjW8X%qb8CJ}$(EaPUfgI^xs2Fw(Y3eEE2f%`=Rsj01fRE+UIKZ>IGA2KMi zfI|5A-n~9h))=XW#DH5Ko?xKW-L)NLt?M^V~+i4iHGM#FZO6;x--}v9++c0zDXl zGIDfpOUsc_^~>uajo}n5M!ln=Xm}(U58vPb1_seRV--mv2K-U`e z>HM~XOoAm7Vrv7633%Ka0Rg^>6GuaGh3^`Cg>a0JsV!Dq)H9-&haoL+ zoj)@xRu`2Y>v0yB;3ML*H|0bfGm?ZtlooleABz`offDKw7S9>+VoQ+N<^H$mJugva*a!066@?IAG4ocYpfgm+ZgKjw`z7tgK@k92V#|5M1b7 zo|8zsm|JX}^C9cP8?cvw0ohYATCfhW1bz*^gT@4DjCA*(^ojvXfw}}&#*Ldded0zq zNkmZR%9^8q&jkf;V63t<*+PNqVGZ7~_t?wYTF%p_wL3D11bh|tK7@?(b=Lu!1v%k&}dl__)%fd-qtobXd@*h zp34|&+3S_f!;hWMNJzXwF2lkfAV4YWg{EUeW@uBvax0kMyPUgyFGXKaBR}n~T@t$w z>yBd6V1U_yNr^uZ(;z@jcMa!wA}hQWn%dhjNb82a4~?-Q!zh9ob7F08-(CWmz+1ql z$CcqPUf>q6^&qgZsbrv8X7LbMMt**03Q2H=xJq+hKkRn2Ed=38WN!xkD5_=ck=8N5)hHW(rIKM3Z|T z6->VFnAe?F07^oi8v5YD6uHqwt-LbfcW&y-n@y{G^CLlhjKHCT5FvNIZtPwU`B(81 z(Hx_z#R#A?occ((HZNfljd9b&`fy{A2mD(w|CR~j3zP{q_$*vollaqSHJFeOEQLU_ zD%rIG+X$vmA=toDn_E<*?bR@NvGFNpoj}BafRz*{3sRWapFRtXrmXx@HBJ7_EGZUL z;wY>uL1%>=4B^INoA}DKnxy2q;M+0g5T1~6202Z!YzPxc=o+Adea4iK#sztL*pG*> z#s|C?=1^g)21~sj_v2AfeJa5o6U6HfElilu1cdJJz{~1`h=Y~H=Ud#lNT2zML)2#Q zdf{{@?2A8f2mVjIMXyIsiG*a~DBvWfH`}$q6CGoL@f3{_-AogO#Y;5&&J{K?uS`5_5+39Jd&Xt;` zWCQ_18!k1lLi!ax-LPI@i^S*8WSWrdcn4<*_03*<&@ z&Fe%3!Vq+0ON-Pd{Pnl*-&^22VRWHS*$11-YrL>(#hPIS+DJMiWE;@xn7Sv^$L>={%1mJ!n{G#eAiL9}=j6}AJA{00poBX}W|p-6P0>|&E3`s>J)0`$#SnPv-ybP(Pcs8kq)t%$7mqu;|m zX?mV-ocGTUHSynp31;8ky;t1b)rrJgN)H}CUj9`tj`#EvD>u#02k0p@MUGTRNjc^sx+Q&YIj z;xU>?I53K}9}zIANKZo}c^L5EZefUL*7UJq8vc2>=!_yIT2B45?ZnUScmJP`*nq~2l}9K*yt?^Mi;Hzob`bNO*xU}`2Xzat z0Z%O^E>1b~;iE^ku$v%1A_?*m2#D7hR30X}0|!28y4DOh21>n_lScN%n}vP^``8?; z6s}lZR11CtW?&%{bTlR}=!c0ATHs_N6JE6OXzWq@a;W8S35msCP(aR>j0VQW#wTl4 zg)^Q#TfvKkHjPDD4X5wI%YotnJh+=p0X1lkYhxX54pBfQhiwb+g67}vJ=mZp^Xq`( zOixb_|H|!uTs?)n8387*Q9ool z08ECxd*J{so^k8*wX z>>GI{dHEm(HqYhxY;2E1y~A#&;k_RINQADf1a+B9(Z;4Em|L-oP6VV3iAWeh z?14ALstfWY5`-sz{il2`keiU4A*76mT=a{*={;-eDD9L{TAZDiR(v$$>8iTH%$tJb z0yV1Fp^hKfE^U;l5Wg~4*Oo!LN*|Xgwr111U{0mX`cX7CYlzri<0Q*j&-UzeT<$<- zn0~w3+8qUPrloOJtJ91p?B) zeDcrXwYZLy{v_3el#0#*P!hwIPx8+FqQVU&v8OH@!xfE|+Z$)W!C?ZV1VkfdIyymo zdvNDS&&jHf!ovk1s^f9vc$=A<1B~T7Q%{9BkM?;4GBE017LY)0u87a_6~q|o%u3(D zKr3D{6t+nt0+WKT!13b?;M|9gIH$vZ3z9x6ijh|H19E{rvM&-SPB-I$GNqrQj^Azp zQ-F5$7v8K2V+2P`U2pFiKr8(zJ@0nyksNEgRo7y7x5K&FQtd`_ z)#i7>C-D@MrX(ypywPa79p=&~G+$9ssEo@Y7?_6-8OkfUa-iRb8tx&F%!5g) zAzE1Mv;nc0gxemHsdVh*K<^c+pN+@F!2vV1pWP*(UvNGsCAv9gir~_&g~r5`VHg}C ziNOMlABh|}a)ZYRT8X742y-Q)hf$%SCtyEAuo%6$*?4chT^Q>qaEGK!0Cu!-L) z8lV^d)HdR~2LHL#$PME2g+im1VQ(g5J&n8E&V_Wj`(E0a$5Fu4GG1hBE}Jf-U@NcM zdVIvz+~v;aL$WUY4~Xxm+nL%8-L!n`jJgGuTr|hp4dTt{&yVrGbJ46CE6w{pnkc`8 zS!|DbyFkJ0Jr?z-jcAxrJf5be%7i8r?)wudDuoA9iPQFOdi^2D=QbzR0;L|@ zyLSxKu17Zj8JIR_Zj2a$K1WNvnZ6xsIy*7BfJJkFv=%QBLX90z)GRMvJci*rx)cod ze*1Q*_;&C{{kyA4(TvB9Z(t(Ua74~cCd2~0lk*`)Zt9yizYPy#8Pw-5j1=N5i?8#} z57=+Dt{4BP?S9C#DgH`*+3^D<2d<3XcI_qEF+Cu1h4!cW^4ECnyKDwg z4$undxU>$A|DlckQ-&=-{Mv-;pzgibc)Oy3f`|?UbR|}Zqq%v2Y=)M~0y9YH;o)xm zN?(aPMqr;H`;v=YID@~;_4-!0j_D{V-F0|n^qaK83JM8*uuN#>VHrZ7dMh~CbQjy< z$XcNIwZ$L2Lb8~NMJo*+0G{Rf^+O2=sx8#i$vt>WjsG~`{8&3ubLj=*4AV7OVd?W& z2d^J?Au3Ka%MRz)==G3{ZutAF@CFCmR7@n_^M{61)M<5$zYn*8cz6T(U>T%O#CvhpX9w8-A`uLlY zqUT=_74LjDi11Lq5V1%vpzA_@^0sX^I-izbD*67Pvx>7K#(R3mR>dS}`dv@e=Pb;{_}yocrT-^tLrR8SZt?) z_`~bvG5SGgi&BG+gIm@536m%FSF~fkStwL}d$T2R?yaM2dNEsTnTcY0aUS=6U|#NjFlJgn=^M(kqwI4PkU1UCwIYkxVF{s5#_e z0d&uRXBfHz!@FFk*7;MGEi{`o2&%*q7Bv9fVlEDNkLb0Z!dX&HK~e3;Cz($@OJC28 z1rFNp^k4Ey^uP=_W)MMM7MX&a+aBys2w5HB=_~`fHq% zXH9~LB16oZ>-4cnXbiFr2!*x4EzJM<)P3)Mm-9hYjy>SzU~=`qr-lHcrlG+mwob^1 ziV|ijP!q=v9(;+T2!;y0Pp?CXwjZ-rhQfhgDb6Etfr+2-No1{ZJsVb3O2$Ka6vd^2 zTOyLS&dQEu72E47TD4~^IW8sTXO3swRp|_E3DV%?xsYSv{>uDa@Di!Wv2sT=Uw&>b zNVaoWO8EW5qkgeJ54B+7-q2h;kpLzd(}vhJfL+pV>g{x^~&X#9AzakjoLB_kGKp_ zTD{SH%v`=)I#cs2)b-rA@b;B7^=L8){qG+!dhi8)vGO43jvT?t1u`jv;wNr(WrDi# zFuabZF8AL+s|_K7ynY^x1!7i2$Z!1K(Xfyc1U4q6y?N{rs5-jI9;3HsI~YVK``!o` z9i3~s-Lt3dTiD5@0rg&w`xX6NYClEi^1_qTtyN2O`ZSpmUmTg)O|`r=w1rlW9Io}+ zJi_J`i~aG@M}zC0^*65VjEas12lOm8m44^WgX4dOKL3*n^;AjN$y|jOjWz~)_xTMa zvOInP8ymb8pG}5m#q&elTyJCxKh+Kmj®F7~lBCb-!HbvYW z1oekW=e&kBNc2U=)b&F@)1~l;o0ITiRDX3XYc`dze7Wee7i8;7c7EvH%N&8{ux<+6?shvH*@4>*!?N%!=>Uszt5di85oiV||54xh4;4IKpZ| zDqWmH#7A;*-X7H0hZE#f^PY`&#hxT|*0{dsqOg5>w;@f>{)$Ov?#Sfe%tOamK_pS=u8HM})%*<*@yF|Ng8?#a^E z*-8VYnlVC=O+-z7m3g6UNi&HJ=O)uEt|e2i8*ju>mLO1TYM%15OkF*bwRc5XjSvp6 z2h@8oCiPIDD*sNaJ)!>h*s$EkB%U?;iBL90a(7#}ulz6;jG|QQ=m55Hc_4OL2(D;bMF>yjsJoFC1mkEcdan5ndLbN!eBPj`gz2 z%Hrq);F*I(Zl0$|FE%=wyrLsU*bR6TKBSd)dpZxk#fk6g{7ORmc@Q8#VGWsu~<^r)oqI>13NUvf&}^a zrlzMwTt@c9J#2bb2Id(wKimtiUg2iT8$9ik-4tsK(?uUi7So@Y-~B(bz5}euv;E#t zv~{-?Q7qy>kgy!!0D&rsgdq?T_9!A%f`Ee|$gozRf((NJ1Ob&W5|)gB1_2A=fFdgq z1p!e}h%#hv{^!N7?f?3f>-xUSTH7Y?`#kqJ_c`a5VQ=}pw|Mbk1C7M`j zS5EfjkEkU?QG=h8>A@oY76$Z>Qrr@08gbJ-Wc4Uo0D_iDD*|BdBA%> z+QKfZw_0VPcu_zS@En7Kf`BcA6c-1(k`xrsQnlCBod)6w$a_ey;{sB*`W?J|_V74IdL&DQ?QCH{GK=JpDFVGokWUrad>zIVv-=qlh&PpFdLZsDczcoP^605zG$4SttlIgVcrDOnIL?DKqo8r;0wx6*aDyAFXnn8$3ind?2#6~H zaooO&>w6cCDVU;bGDOhe5+2? z%f@$}c>W|cw68-me5DDZ?*-J@{}*8hMo9t^AyDk;G`u<^yb@e7)ZNJQaIk<_6mN2+ zgzBTd8&xBcDGQhEbB{eOLoIA}K{fFE@4uIn_#Hnkl@Mu*e+kwVOxY)C%gCc|l$E_5 zikU6e2374tCVm1Pj16d zLmQ0i!-qw|sJLe#vn9Bj$O8VgYo`C3~d54}(G^pVko-J@g$jp9{!2WrS=caNxe&LA7*fDZ`?=xonh2Gx*SAln zOu{I+0oW3DiuD(%l{M_c9EG@*7;Ft{1$TtuVax}^u!e@Zn0Ggy_`RiN9{0W zTgUn|xOb4_J8x@x#_Hy$|KEVAh1vOQm}Q_Fp0$o8%{6ggM>$~00=%KicYq_Xd-n$L z(;)X;_20)6(;id#`w%F|$kt>2 zwTnlKuk&SI>rc64>nPzZGDQS2xlq@pv6Mzn6ZnZ(r4$NS2-1HFOH-{GbYB1E;k>$Y zJ2cP6if`DU_@A%_vNBjPkVZ~@Ea#&q#3Y7y82~#9`%&-ZWp5{lV$xcxu9Vj(;2VLc~oCBS8%8+(o~Fyufh5FpTq0 z)x$svffg^SUeJ9RC<$&cUge4BH=RzMVto&mfh&Oap(P%<`Sqd7^Tde<409NC(T$BQ zT4KBxn4NcBM#YX=4%H_c-Yw6bS=>Lnb@?wS)P{a~c@6mi)SFvk*pBw7tVGVe9{nJ2 zOssS8YWsy5%W z^gB&QMHTzkOD{1Ikm(mjx+(4h08DW|AjtslnOHV5bk5rPHJ%b?2vB@7hss6VmoW;xjvMIg`g#)7|!Ij3i26Ih3-nUyy!!eo}cn{xZieh(wNmCz7+Y>qR17y9gQx<<^Dp-H-($&$Py-|iATZD~fGh1J zg)0b^3#9apXQZ(?1op^%pFY7YN`|wlXGP%&4!9toDq-Ipjxtsg=xL^{9&mHMC^e~w!0l+qI0AJi(Ra(|FQ*X0{({B>gsjwT(aJvOcN(B?=1e) z;2Om4bkzNra13b0$@9uOpwgidy+YIyU6b{7oAu7y{V)$3xGW!~C$*Yr@&Ca(L7 zB_dfHqnn*QLj9Qm@1N@hb30ySl;69=ckPLqxDrS8i2t{-o7YT`0K0Bu{)VWJCGLZG z@F>CvM&Qn!;uw{mKR<>m4aOcaeJlL?k%o|X5$M2<$pqI6<~A61p-VL;Z5 zzV!&?5#Uh3T>(=BQLlwl_gH5ue>WRH`di1_IbbtWE4KNzJU2_4xzt2N1vRhuPHTzO zLO)u-+JAA&(s={`wSlCjb^#c$yuJ;#R@p6_SdIY|yb{6rY63KfDh8sgA8_+yQk-irY4}u6OfEz#E#R4xd)W-=;0PGTe}_1EeL?J zfcJyo6i}?HNbDaf@o{u^hC>`)&v;k;7OV*dIU9!^{14>PFmV-#BwQCvS)LBg41fTqwdkPA7g%gihi~*{q)~8+!X9IWG#1Qu z=8Edxu--IXB|Z81h19jS3)ems?6%fpG*>PtL(^6-n_nn3{Eb;Okc*G(_%9D4J308+yE@q zgXPz{U~MiTFS07cXqfPEa{Y)qqrCP)_`)YxUb*Oe{$_-Ptv7N^A8J#|(YrCV2(6%X3Ay^Csq#q#V_K#d%qdxtel^diHVuI{PBg!*!jb zwH&P|;_XO53DZcAQuOwb*Wpd04L>bVa%v^fE*Nlperu0%u}TRM8(Y33b@ynxs(#pu z%WT~nx^Wqd2eVY?8YAup&hf1z?<`%}X|F{Y<=Z^~j{PTjG7}SfKYzxion6>=f(Js- zt0Ts=S1&niV~BM*oR(48*A-*nrNCgotgM;OC!#xo zK;tKp>vngk@N9BT3Kv`S&^-eFxfm-RS4(T?X|x-+-%`Ozj+_p zbg({*ihJ+P@G2&pg@wJS?dWFduS8}J*NhZZ3k#)R2Q>@zUSw2cngu$w%Wy*lel4S9 zmmIRXq7FrnF_|YN=xTZsZcoqR(91aP7G74bbZMO$G?PoZrfiDGsHU&)V&$Uuz~(oo zTjFUvV{HBJ)4H8p3!V^f9cp{014K{9VA7dC;mNRj8U>(u{9e*lY0 z+DaJ@emik4;_QNOnXuMTVw_ChP<+Uqex2u3PTk`X>tiN?ew9=GqD5X`;!F2-k5-yT zEa$n&_;yrL0ZY;B>+Q~9dn_|~wV$p{&^4XvduGxdV~)*}Jw0LnK)$yo=ZJMiD5n4{p{av=!I$s*TRpg!-OO?p z`brLsJ70gAc*1N^P}x-X&r|8@xyd7s&VLTKNAtHGbMEFH>YO_3jtS~_)Bh3y-QsE$v4PGHaZK8EG*6w9O1sw&fkg*vncMzja! zG(X*~HIp&#+5gF|AuCO~YoY!*Hf=tyP8Zyvf~UP{F;iy;x=(ss701X5;J$VQzw}|w zEorX0_O!bJNsrQEHrgpRM$74K*+dpq>aP}Md2#@z@=O@bxzZjQR;@)^jETecoYaln za!b;n zLOBe^Y*Cd%JF{lr_;ooQnhaYQ$2ZLCUC%yBxxbyOb%10=Otw5_2)xZSl9NK<*H!41 zrP(ejYkGA!a6tgnjg)H1ptGD*>qxXn12WHFk z_x#gNUsSqCkOhK$kgEv6u)8}FKeRbl@VlzcOj*hx=i^P^-9Ab&tinL{DGxghvH(rF zbBhK$!EDr#OuxL1+*4pkw=&>hKv%e3pevD@Cq_e}kV&D&_ocX9X0#3^4ui>wR%TKZeIB&0`_Cr>VuN+JM2&A_yp4HZG%P-+ zWJAg5rm}}_lNuk0W*Vcpke4$8&_acc&y?epOB}1= zJ&NRM{(`bN6xIyfF*N5H3S>cMZ%ZkS^TqlMqd9ClGs0L+5i40)GaSi#D}QgZ)a`&`Q1B$Jn;4B@2lq<3&S|yHxK6ndm1{mBxPsXoii_TS%p#_hNFeL zE*WH8tSMpX)Pgc=hR^TqBn*f-wykD2SpzEUfMOchR-!Lh^g>{hF09J%MlPtrx^a~1>)uM|BpQzkb}Vpc(nhOI*c_u}hME72+33VvHb1$z^}y%CXW z{nos@EbE+hNi@iK3-{M@GvcU<`aQM=oL@ccqFEf^KV-U@3R|uqo0uANO&A;x4?8@Z zW*P2!dz=m!{+$(JA8uYk0sBIcT$7aI8W>7#_@O>|CXom_Uw?e(&v8JmZ`a30p<4iw zeYOWuZqQ6M`-TQbq`2!5I2g^QQM|)<4e1vcRe|iJHW}tt_9Z`xnt+w~4ulk8X~GD3 zx-YO_pl32@e2(UVV6q{$AptPWu&C7Y|V^;6P$q)8N$Ad>$646Ip6=^ zkFZ)Ybp@OWFa|)N@A{KRdeqw{&v-%$*e`e!L2nvYWDUa%%LfR@sh+p zCHT*ybIjt~W7`GT&xjg}czEpVB${`^`z8~^`%<(5o)SXZU1?^E4ENoFGDMxQlX}E! zyCmY@qj8OO7E~7%wpjX^9duNif!oNADtLawQV2Mgz+h{ZCHJgO2oaZfcXw*2=Y(WS zVsNGoEjR03L+Ri*B5Zl>o?}jS9#;`X?l;ozH&VT4h#R+fMje0#s*^cw-rP~^(wIAG z#B0WiKXiX{pY;8b&QF)(WRGmf&7zD;F_9MhlBlK9f)vabv&oJkY>5#vhW&Xde&FjM zXd=ZGUcxKO-kTF{d_NBG9Y32+xt~QTj=?OQe)p4bmtmXe*;=jJj0ZI4%nQ@VX8T}* zgBNbBCd&HoMTFhh!r)um(ZZ8B!mC}s1sI|cgt|`d-}Co!RIK_wE-y z-y;Ib9xBP{Ig&Yd<(XD*ybc9}D97Nq{AQOeuX0mKPmsu3vSv)wjX4K^5)DUT9KSFR zPkCVJ)J%^C+jm#IWS*k3Hu&|D!IV23^&D9pM0#HXl4LU7fMRh~m{E0uQ+4CcO#bOy zT~k^AuJ)z>KF`{~x$ve-CKL9WFzkm>2^bWQTIbbmk)?4$CZdO9?9Mu}lVVr{7;+g> znru1=P!{S##mtTk0h3)wK~Wdest5LEvYc?nTRf-IBsIwa`EpI4GN zWEY&nMpJe(gSxl_@hyt4g8yw-B$uRHw3-=3Rf)3b?Ih&FaROWC{F1@%ZDLV&S*^;G zm*V8v=0A2jM>2~h$?J$#rrU_J@(d^{I}jQ2WjJT2sUFv?PfZ74GiU*D*8hAIXs`l7 ztjAM=tr=u9b5H4+XjQ=!y|3M&K#F}TgYS~yOmRDhTbe}?nPe?-rf8NBvlOM=II1%q zQXEc8v!j7f)+Ae=t7U*$F7b6HRoL3XJbv;9=i2swYl1Z%RnAn}5uk%gI-&|a)s^Wz zB5Z$~HLPC^-#33l8(B>g+?Z`DWMf7O*8y~aEB}3=b0U~JE!YZKnwvup4y`#h+vvPr zK2M#Jq5oxLiMbww1S-*wss21^rkox^HL6lJ@k`>sXOab;*8 zA0>A(_`jqvZy(v5_)PRv5tYr^N2WcFW!8O~|BP|rc5JR>wY_726pX`YZNM4)FtUBv1iwU1B>^tQ@4mQ`MFHpTfx14| zysDgED8(H#Wv6D!dzAi?fq%|ZS09$vcY7L4)<1s~Ey#;CVWTvw8g+^u5navlOef<> zf68E!6BdcC?6Y!LJgz$wtPzja+9VFHRKz)}!dcs$iJdmct)7;jed?JKz_&#gcO zV}WB>bUI^lrib_of87p1edN+eE>5;WSc zKX1-`%~Al|iS20c28Fr>1Jpl?XJLRBhmnjU=0O>9wySz=nc;H@bOujmcbaB2=^au%bo8!PAK_Q)7 zsMOzrks}zN+|1q#i#J6eBi%Iw@iguv635!%U35xV!16-#^RG|njY;lFYEmvdJv(2E zO;P5E-=Q8}s@fY#vTimriEQ4Gyd-8Mfq%xyE?ZpX+)Y*Onjurit2D=Pj+A6yy= zEN<=#|9uQ$8qojc3u5xReP`c8w^rFYYk*NNC`>|7=>RzbB7}L2lE%RllB_J|f+=%b zI}iNw%dHne0niid)T!SAMG)Bb1V(zRtHnsL(ZXZL{YGKw9n@esy%#ES)yXB#MOlSz zOqA%wKZ71D2Cfwo<(Ww^yV@Kb4bfiE(|>icJHw~~rS&MKxQIrQa43BfJ3Y|ouqQ}i z>3y-7=$4n89czD1IccP_m80wL+`fscb(NE~k*lCBn*B-5xMUB>E|R2A+4VSmr%ZaB z$>Ei=lmGq37?uK*)HIJJUDJV!38bIzA;-rMGfOQIqDI?dPeKX@It4M1yBBW)6~|DQ zF@%tya?EabES?vj2gYJkk-X;YgDiQE3z*Fp%j*dWd2TU8Ps=SKT97O=+q)JX^&Lgq zKXN;^?6Ka!HE!L%+i0nmU2Sv9wNL4RVb<$E6pq|!XoP*m*U|4EyOVOp7HsyVO zy~`k$0BBETP`@lFYa&`Qd(K8QvIsxJd;F95dgv-$e`{YlQf4#I=IQd`x8e-LCp{%I zQ)4A&V(;20I1(M}JwpA?Y3wf13cWqTD^Wl1ceJ%3J;ASEc&cE;v-i&owc4E<$`rAU=QndH9jYsC8P9U;rj3jIFTP8BiTjLU~`;b+q?rykugC8>uGER%~_#6`O-1 zxrg_XB=`{UxpmY^^ojW5{#mKqi|*g7{-#^vI=#Yn`^eg*o(`cBKjugmO8L0tI#i_R zdVRjN84h@#*yq~w7gp!}^zDv4+V3a5N51^oy03Uw{^EeAd;0zwFiCWN9`43aSWXhtu{8gs@+HKb&4lXkR*_V<- zlc>$@&zk2>s_;bpxJy#AkA>*J=#ORg*o!Q+Xv=(&;%DW+k52JRE3j)e`yqMeSwm^h zQi>IW-yyFO%apiwb@W|_Jz>OwX@cy3J|n>{RI+^yq6m)aL#Kk~@~R+_BZ<>Ban#a+ zoEa&PN~bi@H1$$-wdu6@6Ow7yP502PHe@?miZ<@vZXjwRx{Eg7mZIUO*Cazu z*<5VC`B0d2mlKugE+=Qr=Vj-ZM~b&)Z&bOEcUq@uxaZDBk3*O-tVza$oG(hC(Ju*Gv6*Y zEgb$mT%kPo2D2ge3e)v?O8a+D7PT7Z{U(~@c(~f4qe(tsMRiimzzObFc5C}46_fTA zY?n_HA!Pd9Rmp{GSgFgvkPonlwLD|GdtCAe_y;h~{I`d7mqdAV(Ut@90l;^#O;>*R z9VX4uc8e~y9m#MM^{A;nTH2+oFJPu<4{oRFP?VOE@!6H!3u&C}o%lN4{k9kU1z(hc_%$%QdET+bH&6I1@z=%&ES1G}4! zY}t~G*_POvHC4D`In*ZB>oZ$RU6T0GXbDY61zXweg&E{daJ)e5%Iy7gZweP}SEJ|b zwMZ*PK8t?C^v~-~xLtbD9#0QxzKMLCzSO*90r@H(TH&Mge2c__jWk|CAN83ud%rdl zpRsVdq0~2+>U5PP&CXIIO3@klSB7HA(vx?OyV*F>TX&7~16KcA#LbY1L{I|P+D^!1 zP{!b_Tue8|eFqgn5@^r*`~~q^AypYGc5; zr^Ru1vU36`ypdK|;5K;)Z;+b1kz28XoiO*LaRc)rS<>El1-bfP`sA~pSMV{NB^FC6 zg-`AL`#ZSLp|CMvi|zBDx(8)?-*+Z75)GdcWNOO>=B0rX25+k!T7B4tgSJji5r|U9lx8^ zl;X(FvZsGqP{!MiQydt_k0x$)U=@-w5kRp&eUGhgob|^22dWE(Ge>sH&1m9gK1%Ol z@V$GSDJQswl(ZtdX0MRrsLG-UEU(yrNQb`#Fjxa|d?~}|&y_nrZl`^EoGU3dU0qwj zhbpwTyZs;0N5G1&5vQ;aT-tZ|eqg=HZF{rnNJLuwNl>w2>H)10M&B?gc%*umcReYo z46FviOc2dCBum(W2Ma{!fZuCZ)%*;Z283y258+;(3;HWJgLW`mDxW3gJ4q%ly1RM@dDyv#_-h>bLo1`ZDwh;) zH;6j*K{Fg*8IF8g2#c7xpRdqU2^F5Sa1{Y9;mOr?gSCg9C8I2}6lG=a!|xqpQnXfJ zLw|r$@#@uOVcHcrFcA#E`^O1FJC3#kvuKp^1I50c*XNc11}|U#cySvR_o95oY_hs# zI%(VL5Ihc4xjkE0S1hxj5p&1JEEMoKL;)~H0eSJscW1C8k&0MLj6B(3F{avOZ-}GT z4~f!L{!$-H|0G3wDMK@FHN!w*kmz>2Fa8IPl#!5SjMi)n?O3p^U5!^HDl^gM59$R$i`@ovA-1wmTxs zU#&~*9CW+?#w*@RMB1N~Dl4ex&J99Zf{FZiXQFm;J=ob0;WY^a%tb_x230TS$8KfQ zg)l33TMed~q?ni&RHh)BjZaK~0?<%9u-?z}Q>yp^L46JroH$SeE_CIHmT)@ayxf&csr2wmA&w4yFEI-bi=lhu?8XegZAqgMmx74^dFUT} ztDnq~!cWEwo26J9q0hmj|fNG5u zM<9Ik=n*E^5Ik%8zI0;Uu0IKnWR(wrSw8i&HtlHf0YQ*c&}`Y44eEy#x-#`1XeQWf zRXz42DVx~X|}jQcVNT;1-$6q`yRChx*q&|{jwS}bjQMSs%;c0h~3~vF-s6A^tbQ~ zNrsf}(>C5dg8Cy0&Qo;J%$YTEw#}P1_p}*Hz1l>R*XdMatw9(LmY{g-pr24<4S4u@ zUsX+X)b-S&O-*o6C{Y6a~AkS|5|9Yf( zbE$FlN)t?}VHyS>-}f;WR&y<U%77Qarl9uPZkeLg{xKB>~s#a5$DgJM-*YDoul9;H(u5ZwF0iy~s43Ia*x6&=J zUm0Z&c%z^(ff>?BoA|AF+`e8UECo2ry?e1fP?Gn#ifDx3~z?W)&ajocfBB3I{`{px#wenpv9Mb#1y_ zu1v8=z+woX1l@;2lcv3lbFhiXEGM-+6J(aBN~5btTxsc)y;1`!OB~t9Vym1fHUv7J zMM;oyWNXDFxZ8RBTHAM~-0M)avPb4b?EcL?y?>(Fk^-lS2(fJd)uLIRZcgp`etMII zeAXgv4%x+k(h@zwe8KY3p;K%uE`|1feBi-jkL~#v$RC^c_l1}~;#D`b7kY4bX1F;x z?1B*~ylr7`fwv0P4#+=-H}}?8SAPPOb<6=T7!MardGPDNW_sr0j#f}tw+i(U455C2 z<^ZNI5PpDiO?VTz02qSW9PVTl<>l}%#kK=H&jzqNz|h*&F$2jO{(HIk`7KaA!V3ls z*jQHBM--w8Ge)v$OXj}xFU?*{KV00qZVoSHjBr=6SP%g||z}r+K z1W6XDRLgEXNcZ3_LSmQ6;e>o=a>ORZzJZJz37?6C{oV}MwGtB3{T*5w``-ReY);y9 zyF2B(n{dujDw|?3nUz`=6H(%yU*F$_D-OP`?c&zUsb|3Wb7v@C8NJ<6XAWOqm=`uR zR2t3;bB9zOpAs_(pDjIxlEbQLnXYepTPC8Y5Hl1@(PkHv;r0K-8)7xZ<%=DoSh;4g zk9Df4vRDB={g?PF@Uo@Y=eEYPo*x_E)gr_0KtUp-$Kb!X#4k+odxEaAh`W&%aJ9L| zfa0No;28HNGs8UVUY#pXbF)*dZkgJUnw;C;mB>7auf6Z1?JyM1cSwDikImrzqpeNdfVu(t=36keWA zmOxJM?%;xT67oJ?1yauAb4H1k1c|nJY;uy=K?D)Si3X5dI+ngK^%<_Jtj{ zB6@=|?6vUtVQ1*$KDw|F85pK?usaIt)k+@h@4sbx1Gb-IV@8ta=qyoIDHJeEF_cWRyAq-WdJ|$Sl z;79Ezdzkx_e`?va=UBe`o)ouI-e)D(Sl&GqvY@Cy$&k|GsnaupTEbI@XnlBFiEab! zklQblDSp#d(IYOQ#wjA~-YguX$LW3Dr;o86%uV7>6>5sl#1rO|#G#&1?>6^Y@ppYL zF^B)A;EV_a!=PmjRIBUaMT4n?Kn{Z|_}EGzN-HVlz#DT4mSoTtg++Z<*A}d*f=B>tZIJ&{CQlqcZWIrE141$U%O)o$+hNCb>((2P z+b{!2-HmM-w=&#kQRq#KjbR6CgGu@q01fFfsPEiAf+^Z7mbOMr?8u=*b$;NBdtqS+ z)DHN!LtO;x8BMGm{H|UHWLovs^73Qw7bdK1jvj5pcIXmQpP#rfZ3fBpVBo{&dZV&( z!QHz?K1jq@9v@8Jpz7WQae$KsYNGOr3P8c|7b#Lkh!wj97r55!C47zEUfcm#Eg$2OnIIZYCBG2WgoQRr~AL*0z6vKa$p{ z&MkM%9CDwq9KpW3xcXnZ&){U={ShwXH%4!KxY68r12?h+W|^ut1H?+~;`EUuFWs`9$XHKL zoh6jYwcLy0FD{`S{n~^^Q8y&(EYEYaNL8bW@Lu&QJb1@ zCax%&?R-3cy1uT@%-Hz*X-WXR22pRnef$3JlT#-$S1nv$o2t5RALzYg|N9^7PAa3O z1XuHnmsk6}GszezLVgAd3U_FykXgagv_~PW9>rI$p3MNJ|q;d7S(Vwv9-mj zIj7^tSFc%PN?6<$IHNO$EmxWsrg?cNLy)nirhLV#fev}+>9rQah7`L4;A^8^keT@z zoNhvHh=>TUSWq2Pyppfn`3-;&?~447ydrF47e5aLBj~5$Fk0sa5W5)j=m?ldUf3tJpZ8=)1LS?Tq6&f9^pc6ht|fSy}}<1-elVp zpBw?Tl{l$=xiIj>W;f*`7FeLFXi1%d3c@yvab$~VjoI>+uz zV#|zwP;4T@la!Vs*!8DBUy!~7$i(mY?_Pw7;$UT)sBm>708?%oG-){R3pk{ zgH;Tw>JKI6!zClTrj6^%&JJZ9Ds%l$i$cT50CGUWpJ?%7r;cAEl1(#!6!xI$>**=M zOm(6IiwodRb7QwlGe%)B&U1Bho2^zV`=iG@7ODN^4q~JRod~3NtGFN6hYy79Lqq7` z>|8DO0GqYWto}o9u^h7s2?LiN>-*|g9*aeW(%$pl1?~#9H}^T%+aLJoNV{#b4SF8< z41+=t7FL6WrlvWz27dRzu0=_5ehJ0IBp**4RV}pVaCY|xG>FmFt;fpW5pK&$N{AQY zsUd0+XE7KUg5mQZ_W0d|qY224%0KZS{Qp+{n)GkYI96j56AkZ?n=Tu!m?lTu`VOn- zRy8R;uc?7zb*~CRbsqZO8Z_`$NPcl_D)EV?q*^a0Iv+ch2VD<@LlCZER0|bK^Q%`0 zkO_zlG@S1~fV+~qY0$8~Kofqg(lXdb+jed2j{ph7drmvyO;Rj*Gl1UZ5l!8btv3DI zFF}zkT26hope#cTb2Kdmz3ly-J~yihYF3c!An*HB#a}BYj!|11N*m(%;d2cDnQP9Fmm%TohPUT>1wogrM5=(8Y3a|R_#m*^oK0zxa#Mi7TJW;y=nkn>@dM$v! zeV|InbM#jcH-KCUZ__VQL-9i(=B4mKT&sK%wvoZBmsDcV4lOEd6d%<#EKkx5c&=v* zJSbz2K1`th2oJ|%uiKL_9t&OB1BWtSlwT+Up&dl->UxKW!SCrSokVTxg^nLXN|^ZB z`H#bIaYvbM1)(Lt_}p=QYTgUJMc>Xc1Pp!)YTgEkMQ9Wq9S=v{hiMZI+aFu8D4jjD z8uoIzyQw+oLMDj4o*-VnZ@L^9S~ek+!=V@M4hv6W<-A1?4OJ67SdAezz)QzA1R)tG zMkz!P`r!ue$jUpi)8CBhr*sT;Ir@|HDIV&(pIoR&(FB1F>%O?18afFOlR!iC~CQc8UHx^ zIQ@aV;CInksie+`rT+l64;Rmg0`ySB zplFo_8y;fB`D}M!r#=1WpD!XED(j2NnV|eamAut7mFTaIHn+5(C-^e27U)Yw#XJ0C zQ@>+)26!$dYUmbFlrP({-qUPND(s4*kKBGcP#6%BXbXxRo}{(4er+q<-H??}8sQEQ z`&vH&E$e-m`wk&NdDpJwO>QsmcEB|R(wZfKmo#+LueBLu>MsX?U;FWbc_G|EQ&$kX zC4<&&*f8@hLbD;g`~H-M<@~V9y1K{-Lx}mqs&#QL zKB`I;f>7-zE&YAkCjd7DjCBMf89+4bG5yH!rfG)U(dF&up19PA#3&L#s!*29r~=9+ zU^O54)?O%VY+Ye8D^DP3DSDIsm1l=VSs^>Q5>or+C0?4Zgzg=T&iz&dA7s z{u!MX?gg}rUBM_w#BWj9$ffC7S$-p*E%1tA)sI@5xqxeFusr=09)#0D_%>=QbT6BR zikJyyAA(`!qen*hXdfXojdO*$w-8srPap~&pym@E^ zbR=0$x62wDaR9reZ2V@X_NlviI<}@8DG`q3sfIT-${3T zQ3lu-Xf6N^5-Y^kZ5j^Iz*P){_Vf>tOyywZl^A;S=eN1=M3d~L0r7bC*@EQhV~w@Q z5i)Sv>T0cr7-rwW6OXE>aT%8uls2VqaBqQR5%rvLPS;ZTGWi1)Y2L|4fkas9pHAtz zt-KG8CBz5>0j(|by5YA58@ab-g`YIKLdRjDr|#%iiBuXZ`qf#F=SE> zS&-GMR$#;Mxf#AGG{OV1tw>9FX+(|n+_^}L1Ee0=S&Hp~u#NNdzo~+~yq0C;*O|yc z&_Ia^!`l_?2H-yi03Qpb*}S#1XqKulPV|&RbF^Jlz*0ZW@imqI$=O?j*tm1J_dW|* z8xCX$2n6tfjJohGGH3O zr6`_>jT(sXc#@(^4g=MNF&)Z0pboJ1yOEeUgEk8Z0UO9`U}xR8%#uPm*a$7W2Dm3a zOSt%3E0iuzkF>bK_~3k}TO9^Ak1BLn(y70c+`r%(neG$za^zMT4UPMah%O z%6fyh@7^I698pi3M7?)d;g%lcTxUEzHDEvpFIBWOkYGXRrY*^usk;Rop3_xjy7(m0 zE&>e9Rf(-ZGEVM8-}Am}zW~N2Hl%)WXQYNpUrxQ`s+Ujfv=OfKxEOTm*s+~huDaH5h)v2k%-+Jn-VRAivsobZgo4&%* z4-b~M%Xly&S=s!p?4vGZ&hZQkw3gTlBP;jos}G;BC$&q_w&@5e(~lh>MHqiXO@PrK z1|s)ypd_M4YOdq2UTsoI_HeSx-KfF@8y~@G2dDR>f{ww!v1cHDb-xi|ALS)GFech8 zimd1kmwf#-RcFes0;>6!V*Wp*#&q_@N@iyVZBZ0u< ztFbY3=%d5K*hRm;Hs0G3@RWYU;K)cS^wJvcMQCI9`%lF-*XiRr5^|(x{r&y0kHsDX zWjNEp-?QqvtUFom)~^Paew{yi-}si%w)FYa!#E~5QMdiW#Z^L7C4czg7%+DShe5z% zgK=ik@;8goxFJUFQ&pwZc3<(8_(kXqQ~TvpC_YfpU`64Zg?q3|IFZ=@@{ql}WCAoW zaXvt^^c!$7gzy{!R5f5KJeit>t&vKm*04zy@@4(ciAxX>PaMX1>R$LW?#$6(|35-P z&Yn4Q+dFv@X!i}Cfo7f6o6a(%GONNZ82a>L`u%kGQgUJ~Yy zRwd^@>ndIP^~wj`zBt~HHkd#st@j_soQ3-!&}aa2D4V#CsEZrpNs5ZmMXMfTG3TZ@ zb%=T5KEf&Ad+8=NCo_|9^Xh5}N!_8LVcT|RF@>X|LLB)?5v8nT)8z@>va`q7U&zsQ<=%VmmL5>s|a7G3!s9M|@B#Y*vt>o!= zy~z+qwnpap_A=Y(D$CiXA18wUEoJ_mIg>~vtSWNGj$(WJFRfnfva6SPhbI2sqy`uP z4;o#a65KVQ)|Hkf1e~xcfYQw@Lx35N=Iqq2gd2cy;XqOP1GauZj_CrcWa0^aXMZqz z1TP$ZaXg&J_@bgwKu-Z!(TA7zlZ{VnD{KYHm1mV}B*0wdkKm$*l6Kdpx|eLl%RlD*Dj21pN_wOH|W zStnOG1$BG$=Qbwcdf{MxXnE0KaMB1J)0(0)DH>4R*AG$5^7HZ@sd=$PU&yD39Y-(( zJ`^_sHT{`iAMI8SgMmXzHcGf;0P{Y~WSIu9Uy*b%0Zs1cnwDI8{?8o-pRdcB0G3uN zllT1&OD?*$HvG}$MMr6#WIz$HjO_dqHB<+@#B3}$r69iWa!REHuJy*XM=O9_AoO^? z+(g3(LT~z73L_AlD}h;La~_$bqyvo1P48;6zf=2lMG!jyxu^KGHb;*LAaQzXG=A;P zpLt?C9MhJX2IvocicCyATK4Wg2@$oX2W(E^b&G)ofPm4iSx&Zqj%_3;aR}dUR9X?( zFNTu>=Mdyh?C?O-(%}NVX2PMYZ`wC8+krLJ-2mQ>?jQ9I>do-25p-^x^mdq@G-K%{ zs^8F0QYP<^KLA}KlTl(owDf@M7TPCfOiZY(jj?e(fD`tgZ%hk+05AY@^znm9Y=@!i zz`V2%A4VAl1>Tk7sE~s4=p+o&hpwYxMDo^|tO&IpUhx10HTG@9Ra1yk2Cg5vTDMNr z7afc?X!xnYmw*-?A9jKMXh`vVb<|<{&o+1VvFKB@wI4g%ooi^w#HFl}P(AKbh3)DzlAPq3qpRBN!I-&#dxaF&{wC0@N2#ICGX zeUzR>rdwTAEt+1zK1RGtyK(%2QI$aeo8=9qJYwvOU(N?hYb`NOWf)L~AEhhuw)(1} zF_KV}cmD0}!<-qVwj=+aN^Rj+BpmBb;2uNAu0j|aQgrNkJ#!w%Vh$Dl!ki z0^O?JV-yJzm+n%^a8D#^LDkHyWai`ax4cpw(HDw3@D*PS3nY$WNV}`_*mAcUBPQcJ z|1&jWBfTOJ5T>!QvGSG$f6S2F>KB4+0bT*p0*t_9G&D*>NfiMrhT&so59b#MBU413B0lZb-LY;psXnTb+d$0&X{3JlOLBq->Rz?|(yVH?N4b}*LoIkK6- z(siZ;7iA#3Fq-qUI4o(E0|2RJqy*%qjrCC}fw`BmGWXeNT$!EY`FsT{CFh)Vlu+y< z=YwaFlbLzF=;F#d^(ProQ5a&u8E1@MGB0U)n5mVZZc2*%%g0W$r2!>so==#O24cVO zw*7Xj*Y{!2;QUX?U298cZuXjcd9Iakt_^QH`-NtCt+ztKo_?qp^uJ?N$Ug5DxV5>f z&kb{av>5ruh7qUE?0ha7zE9q3cXdML*{;}UXMeqWGsk?lW@>)s>1dTOM+1Sx4HNSZ z$D{XManwJhE`4|Ci@e{>(^my4x1Ai!n@>jTKG_zn%18!Xxk#YYCpPBSjSh|@&WdHA120sGk7O>lnBEGVKjsSjA7w%rbDz}v_>$fm+F4!tq zZpc`YT~c*#NLfWluJUfZ_E*F<)eS{I#U!P!F;nT%sHouIsq~2+nQpv9J>%3KF2Wwt z>L8g8F3u=<<@5R%m8mhUz1g3_>0j=>@Ua#RZ?}q&b2&mQIBo5!u9&NH~K5gLk;Uo+Iy+9@>a2gitkpJ>G`N;wy$y80PnZC=aoPRh0 zmmrFt(=5X_LA9Wo3;I`-dADn-)0(QEXUl?9ha;!C`C4)}n1eVIlXBeU%imnMFiEhjf`T|ylJ`Hmz`t5ZDbZle2J-f0 zcAbe%Q|R;qzrrm7dhBMiB732oWwTiTgT!b38kGCj4tWRV>Ln+}tKE;ZHN0^6t~zX0 zcx8F+ADxQF_qofj3U@Of@$DFkbF6iEb76+dgw z*_o28EQNmo!mlbVJ-8ym^rRd!XV`kc-ULZ%U`>(;*D?$8juW#>!nBhqY>i)PY{6x6 z8t#7Lp!_gE(W2wrn0?gGlJZJQhDbbyhO?iG&aE2z=|EX2H z473^_3@{XW)YxA~VONbS%!49Xzg|7IZi|Jg#1?a(ZF(#1o&GHU{!zmM3;H}giEK4W zQ10r0=(lNV3_pv!m3kw;l5ORmnRrS0+&$zK-Ic2LgSEw_rFuihI#0aU>gxCp*h0l> zdF&2?Ni?w$CPFd^m6}`qpvs;-D3AbGhWkg_P42SHiq@cEx`DM@*pf1lKS>5}tJ_#MQdD1qn~5of#+wf$X2^%rNT6`>S+C|PU&DG( z`7zL64xkILbeLkWkfBjTR9g+5RInHBzb*ciRtN%tf_6fKe=uk@$dcX%z@w26tVvEK zALFBVLQH+KR^`f*^!?tnuG#f6zH+oMu7G62S=t~f#S6q7He9d2nXs2MHUbqRAN}Xll14U=Kk#}0j|MPDi8b^ zs`nL$E(650Z1Jy@wG=WvQi|LOx+Vi8Ops)YRkVSr7OzUyW$;JC561dfYFys1;-4q- zH(&}+Mi9raWi~OGjqz4YQCA0*B*4@PHeZ0IXBxUtLcqcE`0@APgurYOE^`2(n=&mL zr(GNYzM*1zi#0p%-Vqib=&nG%_~Mqmdny7O?%=Cgh$%&Qc(mQ*g)W6hs5TZYf-&3U zM~|K=d4;zfM70MJ1->pu+Nb41(9W=okOKt%zn{AJJ6Cko21s`Gagb|do6K3;* zlgb0y>Qud0)gyUid|lMX^hxB_V5BGJDMl42$z40i223e7V*2n&yC`7;gC)Ipn6cdh zMmR|WvmJkv2>QFfg2tF-x86Hb0buAHsG_=^U~GT~fHX4Gu=(+lz_(~#=U|p!aXbLb zA0?Az6%=~>Q!A^t5_a^!ePA+#%7R$0ST$*fWq;71)`P4f z(TR=R(UB3z#E3D$i6}cGBdn&JtrW9^|2bBC8~NWqaVJOD!s6AP!=czE2Jf)(`PEwj z?zr323tfx&(QDBfc+e*eFyLs1-uNn7N;4p77(5Vmq7IhN`L@?wxFkVR{i_@-2Uj!dkG^fz5?JUKo z2!=i=*MNUs{eNt|cOcj68$RxwPE+cXL`qJ>UfDvals(EGp^R)w_G}YLg}0fJy%Mrl zg$f~iQz9#rz5T9N>ht}5&-wOOpU=_T`}KOB=YH<{y07cH_X|GsQ{bX?IeYdmH(y|u z7-g)pFmTZ4>lzN=?M)PuFzKnQ13NC-F*sCEc6ZDOo77m|OY5(}_8XHtEKT+3L8gv) zqqjj*W{d1reQ9_0o2{`qKz&Y}oI7TR5h;_^zb~~YLTq?L7xSEu4EI)4#NiBMULNM= z;!v346rgT`$_U1xcmcrNVEzX>7NQvAc($}9-YA$naAlf-VGW9-%uF6`{@4^g4p>uR ztPARN!Y{}v1x0r8u1l?em62&1T)83uz-@v5^EF(gP0er$#GIBX6xA4`letAX2BYqp z;92vr8jY701V+$;Dk9K|0>uv$0trdU#w=KHCQU)$gXUX9*fPsv6bgEL1x14I<2pUG z>jcpvvQ->T93Di$0-p%Pb!ATgQ6~CS+wS%&aPZ1)h|WbYKSP0^1X?ID=dm=);ro!$ zzvr<<@Ukn0`xKvczt>X@R20L7g*&ZCbqD$JO#m76f(|!%+>%RX;_26R*miexgq9=% z1SU@B6g@M+MCi5pRJ>&>S1PscfqYlqSHtJf7TySX!NZw z0U|_Rem-DLE<8Rbr$Wy&$qOmXgl}7YJ#gmBmoEc}MbB&%wFzjP2lnhK&lPFjR(JuW zjspfe!YUEP4mMF&hJhsdu3Q)5#QgpJQO0ce^G`S;qxG7FR|<+4_T35qcKzB`x4^(a z00IJnf|xy$ZWP2vEAQlVyxSfS1lApK7*`f-_$%;l%DYuf&tr}7E>Ssx(X)SK7Pka0 zImZ$WF#cfIj=ssqNASNT_-R`y89!#BI%BQ5xVESJk(Cv&1yDC3HEj6sQ4caZup(Ar z;OS1km<4Z1ghT{jOovO@)MV-0f+qz|v*2cO9ugI>$DZx`YJHOyO0T!IwGYSl5#vrl z8aby9Iz>f-`Q5P#1fKyEupbkLo?`rPLP$nV4)Q76cQ-$VV3wqFH(&AerqZ=l7c+@0 z=xk%^TMaNHZ+hUk{kP!d(C^*oFU!cv#-zI$&%BO2a4`OB?NHG_;FTI>aFb z{_1OJ2(*6&A8Opa$r_JjBxRV{+TPGo2Hz`ruBBt#l%0C=p%GWhK>YuKA)+XWzdN%W zjK{1t0NR58{PPKPqRG7{Jk{C^Y~WsurUobyfU=-SB;x%`fOnY2^F-1C9@Ibqz-VFC z!BoaKI3;lXOk6>;@4Agm;%V43!sQQl6K*KL7ve$RHX@h-#K(>e5Sy=qg~517Hr&be z(WlOo#bH4Tvj6CoojVgYVWZ$=0LKAh(`N~(2X+ZaqKS2YKSIPjPEz=GeB22#ketS2 zM_{{pu0J;V21N9M2{$$-KL(CV9_*}XtqUdXIeB^g-%5#tfW3XkjtC?IsEUE54{0I- ze6f3uGwwP&^c9orG&2pi`h^eI5pBdGV##r0r(a{soL5icMe}CmFW_6?8X|ZgAWsJ% z_wo+O!02BLK!X7ua?%sGwsHj&@WI_pM5GYzXHw1`XK>NHecNyX=RwYzXjj3jx1?^; z{)DEgYBR2rAz6?F)E4+JZyZKK(_q*@9ks-@sMoSP-o2R=3nyY3z)8)q_9eJTU?eft z3C#?uZzcbccS>4-qMBpTd#PM4xTRv}TO|9SAfPOhy3)T%_wRVXkIL@*eh1QDdMfDl z0r>z&bl-G284c{o{;j?hfS~bGiA`JzmUbf&9?VJsvhAwV4mgQLIUYeDBqSb;k>%o= zAaK8AQHOD);?V-QmN33b7SsPQITbUuV9fqE9i1E24UL>p9g7@uoZq#DI*+0I=(j$O zNYI9O09~#p5uoXfpf!lx3=kfM(-B6`Xza*2o(6sQD&}nEkiX-i&vWz@CT^loD^lFT^=S*aTM+E$|^VwBQ;%?QcrRUq2W5tnMczcV>1>yWlen zd0-=t3RKDxbPqN#oJz*QbHzn@?=k4onO2aOyBrSr#Ir% z?<}TvCI1ykWHQ9id!ef3=YoiyOV^*C!*WHxkO|epP2wO+G7ebj7bp?VVuo}+!JGM z5XJ5-1A_(_13tAyZwX#48+Pb*10#}1yLx`-P?djxN(dp3z z6;z3zFEL1ziG!0;K1TSYrCFl>nU2=^gdWl%ad@e|uKP&TE=0rXPl8&U>+9naQ@fo` zcpBGX5rzB`kRkL23CFLILw28V9Zl~wxfglD<*m7DQVl5erkjSEzjItLO*Z`m!XF@V zX^{adJ#33moOYqdTqboU}7Lw(}tk_nXSf`AaD8vc5n zxz;XDk04c5c1u@la%E}`t)r#ef3Mgt{b-Y!orOgzn&`%V$CEJh0t*1gWw=$k%4@x8 zYy>W>5XuLRX(HlOke>VJ0Y5}RZ7`5PCqqquD#89++XY=+&K{U6njwliw|-qoLHoje z{_|#3y3XO4?r{YbB|J|tX)C^V${!gwVSh4n?C8;0gx|b8N0`h3%pn|4;-j1DKtB9f z3Dpo=pv=MehF9uoS{3`xSK|zWT1(r3IeYIFSD}r-Jg<1iOvqhOy5Z{J%WcX95f77Q z=NA@W6fwD}=_Wzb1|3o2G1~3Cij=aqLq!~(&+kpFL5CHD5BPVAD}HWgO23N98*v!? zm6C#aU^H%3ot@fUc|oDD@K^RgnulOjrZ{p=?+Y4DBV@($+>&deAHxZw(T^Wre7RRB zCskw0NQ8_Er45a*J^ZNrA0(5-SAvqBo}QL=(dkYSM);vLN1Ez({W`B0 zwh7|l*r!jg%9?N_KuOO)wS;$0BXl(i7aqunUHRAHT`3b3>wrnECotz2-8sOu6?>-2 zz8}l}!c*$Xr3#HBW`wrDM&`@=QFF{gxOMyXuptsIq~N{TD$DsB_OPL`B1bp8l$~@LvD(2=f zRCv~nU}_^7!>MO4*$Al%IeARNvCF}5DtLAoH2@5Xv5JJu3!eY?l9ZC>6B0XZd%C;3 zTUuuODtwk>bzNWg!j$jwWgXHaS-s<0Lu)nN`i=?HY*Ou%K< zg1cwwn5Z1Fz%T%v7Uep_&&x|7CfSdXAP5}9Fvt@a{%MP(l#iLaF;0xIO-J|y_fEfJ z8;CCS?Ims^D|X=XFEchV!q^mY_@*_;@a%cFkJH3A5Xu8gj<^Ie*p$T76o#r1{f}!n^9<zrW3$XN1vaAy2MO!iB8iS77E7H{gciBxfmX0-t#Q7Q5 z6wI$($8sVMVf}YFEQUQgwrd9(SWu+Fupf%F7)<>@y@}V3Hv6~KC1!uvIU-mg$`ySGpni{QeNzAY zVgT|g47ey}r)AmMhKLK@3Qlaf(H+x^Xlq=*b`46Em>)7O%^O4MVSa+U3Y<2u;p^oA zJ6d<;oqh@$NeSk?ryCHGz*1D7{&muz3E2gvAD1hRyC4Rl_o75@!PU=y${&&cvpd9k zd3bb`i2tiX1uwxXLVU*k(^G=jIz9T&f3^rlLBW;1eH}vt&rg@2hi?(N<_;nhaSxHevco~M!YEI4QTc>1 zxe+6=;*=H;w4$Q^LB@`a5z>(m8+)<9))Z_U&^eK_e$F>Nefl+VLcg_L*4GDrVG?3Q zqznX1`K*K!g+GoPel^C?1RR`hR@VbPzdr5xRwSG7_$&{~0LQiuS2D^9I!F;v&m-Z) zC^V?p2o`>?XTjLO0Kdt{D7^3s^oi{NFH9crZ5yGYB$2PKE=>%t5*=TNcoES3K59Tf zPE1R9{P@tsMW)0ukh5LeKf|>jVA*ZwgETZ!b0LD(2lY0E6oB?-;k=0J1tqz{O0<)` zt!?Z=?BfP6$It8@i{K6HgeeVj7^I>U2SxKx?Lr=n(axXzVw@iO`O(8KiCWLK9`+F6 z6yV!uk3M=qd#w9r&#l~C!D4pJW75|CCB@S$&I!}IO-UvrI}VCU8Dp~^QDeHUe{2|& zy*CwjP`RBj5kjH}-Z5XNv(IVpsle8hZUp2hs>^|X6QqZ(t%B!SZ_7gsI{uGuCzx@!0nxOC9^TT zAGFc_px;R3DZ&cpZSx3OJFo~fCGDTVVHu}$ z+SdFJySIR$Q%1v-3Pz)3=q9Qi`hq2cvWW zv(co8?_87R2r)S^Qq<{$A(>1OdI z3}>5;A2OHg#+IB0q7uZM@6Q>u0?z?iWsuU2U0eBF>Mab5^j*$@e8ngZNYBYRp0%D% z_Y)RRf8surM|+HGaA0lo2D>Fxqs~4N&1QD6_ATD!+=59tL;RK6TRvOiqcNirfxeYH zdnMLXHTCnx>ZdV~0%tx???F@1vx$V3n*_KTGX5&mPDn(;zMg$C(kb^io!mAsE|tYs zRU+#F3wyM9^&6>fiR4JS#cqJDAh~`ca;263ZWNFfQEG_!9Shvb)`jBp3g8Uj3WWLq zUjvEAHM}6p&YT||Cw(RSgMuUnor)we=i_T#qzBlBC66p1_{C*I;O_tIqM-rN^uhQ9 z3zYXjvGWPA8D$ujC2c zKTai52pI}jH@A1b`*XpG&zQ_a(U1`74(BO9?5o1U!mvv~;gR7br!}8DZI)f|;srss z9Q+Fj$28N4X%KmF`wDxkx}qjQKAxPsi1BeK0PMii#TRaL#D6X@s@}V(=5_xcyF84P z1MtRI+(bE|h*9cb1S4$5-v5UpvX)T(h1ScPoF&F$A*ABx8XFqQugtu8n1c-$mwR<> z)vb7AC{9w-K9JlYs0Tg}NFq(vo;O%@0%fI-{28@OOnf?b4fGe+hZwvCqu=svASy_1 z?)zR|9}D%bI((FCuz(2kh806`X~3E2s2~xmVqpSBLZ014lww)WHk+E z2#SvL@H8N6ZET#vfpli_+^Xld$SJ71qGB2nPf`uPr0L%QXDwI)-D)CRL5w-tEGr}m z1jf!h_hmY#gk_{(&VgF?kkOu3A zGCy0T7rOcil_)TC08_a75VzBllXvgOJgT#jlDLvJx3DunT_A)$5Kz4-Ro6nl#ZxDu zSBKJSzwOA@dKw;%;Uq#O+GfUoFmZi^nAoS)ou-=bemw=7+89Zn2`$&C338qNXxea(Xh7ETSbK4di01 z_R!I!W{>wSOZk77schJA8-u(4jZsXf{+#}^1OF8eY&>SWbCofr!_xvN9Fgb}m)Z+; zi8QswFaoR~LGOVRfzyN>4f7}(T-I<`dMpkn;QNhquTIfd1bxVcW*nUK`}ZH~k)lom zEp!8=3hX9i#j?`Uk8_81bFG1AVvM%W&<^y|&@T?gMTN})|5*zlV5kYGsa;nOmX6?o9)g(*q zkw_5!MPgS=^;H%fZOE#1t41Pt$BebMu|K+(g!$f&^nQOsGPj@RKf0pc7`SmALJE*- zLVG>|)IRB0-|hEc5pbb23QA}*0sF4TjE=+D0FyKzc_HXzk?xvi>9Izk5^8G!R906) z{A~l1IJlF=WJ!BsbG|)deVYj8_}-XR-=qAm?=v`$kY+iKm+x7K95|Vzq#rD^W+LkN z{nWn^OBY;`;UW$@E}PXguh%#j76_ne#;ZytO-)-#Dbsbq1{X2(Qp?RUjXR_2O#l}l z%UG*#+Piv?UwLmvvMjscEhD*0Nw3cYkHm;PDblp@Br1NG8MvD~fw|YDu_>j-s<7IU zm#NrH^1vFoX)w|w@&D{l;tMQH8bM5qs}+id1*7?X)9el~4zOcN9Dkf~A5$hYX%G*B zu7z5qD*beNyiw>THMc-2D}7B~%@HY`z+_HEVD+x1y)88SuZOu9lWGgL^Vj>t-N{T% z*ZlEEOYr-bv`&&I^`^Nbw6xH{+W&8+gA4p0vGAZLD})PipE#l8zR!6}+9~&0<#Re| z7+KH4av?=*u!iNccv47mIDd(;(L0SES`xocoS#AEP4-I3pag=}So%+(7GuHi$%x(0 zv}*;I_LEBQs_GbcU+msiR{TIB<~}k&P=HAb|FTXAIui|OJw7hZ2oz&n2XNkbRVq=< zH#XeE!~~BR^?S`3Jk*9@)q@ znba9?q|}`5Hy0a@RMokIHDWg*vCG1p<#DvBJiHg+1-td${U1sHKPrL~6@dkwG~9Iw{w(NcSHn`**n zfJswlyLoNX#>A0(J>k*bL%qwOIlBL-pawk>TN`FwzY&FBM?xl!2g9WiL27Eo&ey*U+c?&rJ}dl?Jsps@DS$2?J>x2s#~!(j_1u0J$T)Ju1^zEHRl4tWGOk zb#5@}Wpu<(N zJlU(Jp^^BWM=%PT4>)Bhk_rkJi%cXuzrsIvdF`AYvHpyYH?SQtcx`q19sJ~8K` zT`)57i)TJ@!j0(g+}pmuWah-n!$VAeqoczHBAoHQ3nn&?Kv-Ya)fMiU1U!P2`}_6& z3OQy_e&8$d9y=DyXZi#*6}@}2j={oqBPyUBfgSmtYt0gEm)cuUJXHVufaqqPb(X-6_mk!`}c!Si9SRnm$(b9FCzzsCFntzBm;;G1;6i~ zaY{)Jj}6TL4yowrDOwI0FBtvd_J}h{g>~@|KY!r(w~>*mpoQjb{q!a)Ud!xcrPd= zXpEA^KivxtR7y%^6_ukDLPVWG>JGq*79nv)YJ0)-&ST_#ScJ}dFGrDfW*DA_6FA-(;J4sX$1Kw|U$y_b95 z^}&ZlyMBMsA;JMaN*yQ&Sd@@A6D|!%>oOXwTcOS`FE>O^_;co#6&No)Jw1py^0KmD zC=-9OlkxpTF2DH;tRPG|!ZgdpCg20*qN-_X_#Xf?z`m%#gstA*NLuKaF>wd&MdBy- z?zKWTi!u?}3TZYP?m|2~PiQd*jQX=6nXM$GuPDN>m%e=aCbxDtBxJ>}ciTUEXv4bx($bLDN}Em;3y)q_YOr7c%C0;-oh7bGOjAO!)S+ zFI(9k6s%_dCp5W8Z*ZG2bgbP^yYr8Gvdulv#=CQ6SbrCLJb9?dstVCVcE=l#h`9z) z5l)MLF9jRU3PeXE9oZ?EUpzZYFwNUc2^cp$J-rQdp&hQ@Ki$;e0FsBN0<=ImIz8HI z4MhjwA!;^0vnFQrro_a04S66B!R;}C7mX$;kahF=XXhZrK^Vb&wEY6aL^GVILfhfK z@o~7kXvI`D!v0}(6g z2@G6KRZVFDAXOoZnr#&2!Fo4gqATA&B!tD@0xzYzs|yGP_)-y(k!y{LV*HjJDyh6k zU(j&}v~1Oxlct)Y2EQK2hV;WAxsKp5^43kKCI{HSTuk=kvKywD>eFRB+DSfftW!ZH zCpJnYGtiPAO<^V&HUOC9ddrLlMx$0pLL(11Zu zKuszq?gicoG9qN7ZKkBVXkw&jXvE0`BB8<`?8W;}H9;D|Vo8lcyo2nE*7q%NM0CNx zrnH25E43GrZKPzh)%0aDo}C7KbM&jZ7iq}jCHw>8)sTBdxH_9;np@T9ql>^{A57=x`! zYzAbqwYpKS5V5RveVqHLYs|`jL=^$z0eiv~8}6cx4u~}<&M?D_i%L6w8}T9$INENa zg$S-wdx2d_UU#fCWk%-9TWu-GQ{jP-q*uBM0t|u%v1KuK7GD-I&6Li+(iU>EjE9a< zwgsO6P|S>Vs#}My;(|;+kJ?V$qTs5_fB3j}n(i;~Aw!)9x^>WaXp-HL$L9 zm&vSpm6_iNJJ`1A8FtN+QU$L(mBzD>C{5sPA!A6Q0|LVR)F((EOO3L@Y+*f@&v-6EP@i z9XYFGZsa~x(U2L4i#wIw6i%9LD&VVm`2^V64vOJ~qh%i!BO4S)B`aywl614j)#4%> zQk;#wCHxmAKj>diRwmxPf?a-tk;9%!>fDnl&fmVrOS! zk}Q!pZom^#t^ncvIj9|Qg&H|!7$s{EoivbOT%4S6junNq)BXGuF-;EcD2LPhNhDM< z!aH~RVN2q2^!6ayi!jhZZ3Dd6zhzhC8N{kKDQw}`>V4h`#12BRiStRrPWIl9xxqp( zN16Vc{Ku_|GXI)#iuU4Hr8ncm13Q#Q^AlQ%(eP!>j%_(-s&yghT;(`b>V7qI(mho( z)wsyxU#B+6@~ZXvs;L$Q!S#eyHQDQ?5zhXPA3qTB^=%7qWBlJwJp@YyG|jKt+e3<2 zXmrvjA2+KV+%^zno?{vwLJ4Vr#M6wyY^VvKR6#8RP*DR7u~Vl`Jr%Y?JjA)dem=No z&ruSBljO3p=m?T37$RJ?vl|1n3NDwfZiEKmqogm=mw}(%OI|?&EE3x~3bPCcOeGxN zhtcm2-`YMSKN=WlesNj+8qrJ1u`>l+hp$2k)&VM&UolVpq(It~HQ7iekx#m*;8dWn zN|IoHJ(PX%H-KU2hR~i*`OvWk+WPxcaI!(ru|5&DcV6((*H{^&WDA94BduhW7#VhcGwHD;zMX|fTgqM(WE(+y&)3ED&k|e&*%kCF9myXO9s#?#m=de36j&qA}mw+gpeVUk9ItEQMUvbPkkUd} z1Jqr-L|{EC$VlLJnn%QK0!2>NH~@RyOh(4r)z{|8i4!Ym8W2>qmnM$~LgRO5Xyc+Q z#el;XFJ9o*2hrLYHZ^^3F}KXDJI~O_2s9aYI8M?5BST`3T>ys8bKbcwscvP*xU$-33g8G|+;bSm?cRmr&g&{TaMms5W^7K#7sujNAYno>I#9*y}8qTrHoS+ zGYI*kwgm*33HzPxC$S~$u3ZZ`ew}eGYrRZ-Wgd+^u&jud2_9v!k|`23omRHlSWmGC zFRyaRZe2-tcXv#FhbRW<*1%k#@BDduU*LmCdrkt1%HUt1eWiVd6)!&dQfGopU~-&l zHm6z~PAH2^+GvA5b$m(c1OpTp8$YC2MBID59hYj9l4R17Z54E?&#FcT$TT-3S82#7Yubfgxqeq=bcbs4g*c|{@QzmJGs(+w zr3~du+t3==vIP)9FoFh<0nnZLjB+sCGd^ye@>|47#E)Hwiov2znSgs}wj&@QKV05l zmt0OC5ER5T>wzmAd}k)P;75Uh40;V<^MMLa;SS^471sm;yNX`*<-%g=%5zCuR z_jBidnz?`xOy2LStCw(aOXmB?eI*R|Xdx?gz}W=_a>03i6it7GckstKgoNhNmk#Kq zC|Y43h#MZDQstOQo_U}uyTOb_)5pjy$jE4PM8!;zqdU!F_ef`Fv`l#+<%V2Zo|mS|X=eJIB6g)~9tfOj{nv8*@%(MFK3MKg`viBrTGz627 zM7<@7c3I<%y#NQ1DkHtCpg9N;7W@b6j(4u_jh-#o8c{UJxj=Zqecoik2NM_I?C@y< zXZOOv*3PB<$QLnE+#J^(k!bgB4%nd9sOF(wrL#9Kb!w@lmkRZQe(K*x>>zD?KSCs83^Iy}VmAF!*!GCBXQDrj&fsV)LJCjFOmUT6J)}xDp8x_FGGtw}D)5NN z@%6pbSOalK_2s<2WYQ2@iT(!KLx(+R*4Imyl6Yr)Xn&!+|AIgTeF8tNuToOe=nUAt za9IO2-yJhyK_h};mNL^HpC%rIc#ez#uMveQrb3{cyc?2tnU5AXJCnpW1P0zi>(fF$ zDV|0|b-as}2CJqfQcMhb;^t=6I3<=q83>+n3GKuTZlspPZ%5Fj7jq9s*ETGyTJQTB zg&d>u>(Tcf;0rrVo=ZP`@Nmh&c8^m4RTm>!aXD-hD7V`6F4NUGP*XzU-m*Vi`+nEQ zhW&a2K|OJ6W#UCauGepdhKDaBa07uLoVvi=0pYmej!LmEq!pCl%9%AK0_g*0RVlc> ziEap9-7HL$Cia>96750M0&)k>br?I ziB^QuIh=s60XK_WG7jZLYHr$>{zNy&Oui@yzv%?WZBL`?$W984t|REfUZjhCLq-Sk zeoj(55+AhaSFVR(^Nb)>c&Dd_HdwtNA+iAIL}X;0k!LN;a1F!sdH#=cX&3WtkXTq3 zrC@7voMc}ZR1Zhwg5GN(H#QONSg(n&Ww!YS^!f>T!2)&1H*`-J=(zN6JfYXYzbSN3 zbiazq?k)Gi&6Vk4Jdt4;+G^VVzG4l>OymVV-v4o)mL~2rSL5;;g%=aDg2#@QNH8g_ z31A--+^}Pp&!G0Kc$Ly@MzQJkdO9Gs&`bH*(=q-YFo%-$$&=;BRtO`i^%Y(BpQ_*yJmmpD0 zZ(?bA5Y%Hzwk9(Lw*pfQSfpq$EIjKrAuj~)93Z!%yaGt(3h1W?eYUwmBO@$@*CGVRm9S?;@D&ossc1=hX8D(f8slZu5zXWGK6fy_;sNXM<*vdtR z@Df8ph(0bB8Ntm1{iaCwMN>rRTeofzbyT&e$Av_+@Uh`WH6D*=a>^#Fr*oQ5qQ|Y4 z&5Vt#xoZfS46ZA+0Avw+Fk1Ur>u8z}#cAa~?e-9J0SQ`a_8&=TYoqC*QP+sNpcm=2ok%>WM&@zLV^nu!Wp@ee;XLr!&I(m$7_;2e2 zhIoLL;ITFoT{VYsh8CEiwzZdk7*Y1Vg(}UD%MA4r5T0eg!B}*bl78$5OQUlK{_OL5 zv`szfdiIGZ877f}ftVX1GanC8unRw}<(c%#kfIGhn#B=AGCJkY*zFyyNV~m4v08{l zLdw))k032hbb*~Y_e<}&LP*0h>VQ&^=W+7?y1l({dw*v_-UPyx3rGzJ!iXadfHj~6 zqR9e{BQ`eHp9_dDtbBC~WOq|h5s}zU4IUzepdW&Hla`%{)Bur$;{^|v1(20u4lFKW zxnVNJVusH|W;kX#0JE`itij+jKmh~`E+-L=wEXx7niq1i#*xneB;lG6z%cKKOH4_~ z{84MyPX+HqvEU!M*5s4cR81i{mXea* z^5|Q`LuX@<+&A$I>(0Pr?96W6IGi~w(Hgv1mr9SG9Ktza`7pImzGJw%Fn55{2b2$e zl->qw7NC6vAHV+fn^-3ygz??a>C4HBb z1t~ZC@U1n3#)-WA+w(r(T+<{+n%QbEi=m|F^x9qx74FEV8v4+=ubj#_T^-HzP_lZ; zTV4OSBeDhK3Vq60n`3ndRt8|!cjMw z&yYA`xat~iFy0g5{s>v$F>u*ww^jtRP&^^+QtFxq_RWRei>q~d+Sb4FD819K^V&+B z*QoQq3YTX%GPV#njH&^=w})~?SpXv1{>EcRjunO!XRu3c-RFD5G^K5XRB+d zcLUZq{GB@HH8P;M*Pqdqz8Jsf=ZT3DC9aQ8)NDKSgj)i_iuSk2 z9$+1^*!GsVb?TNF@_<$P?gqVx@UIw@xr(*`4%WO9*bj6Zgip6of;&eGHW~*5A9eO*+#q30NHB)g zU!~!b^xu6kBLW@4w}&bcKFvt3RG9Z5o5qq8TG9X4kKx{co#zGSJ=n$wmNnJY{{Q(g z9oT9l+7mZ&`KtEm~;+ju&uw_y68j|N9gGi9nX0pJ&F_L+B#D;~a$3hQW765cr70`2-YCOdy60943N-H$;Rf@m>G4sDlNz-*}#9LKVBe+!x8 zer{rG2LBL?%OGNdower*kT#;BPJ)INfn^To7->24{|Nic`+n}lYNGCpjqO4zi(3wh z+Twspn?p16EUPm5JnnD=%s3_{wgzrDks=|C6RNY$ZNDcMc@tF^Sq~2sTDS-?acdZk zLb#1nCs6k45&;Wfz5T0H#KcO6hDjJ{)EFrN(uV}}m|BSVxE!0S{aG#h5E&(rEnkQ{ixGBJhYS422vN>u-^ zItP&?fldg$F0L`c0-U`LX99f#m3U`oXEYMvqF{mi13C`a+Clq^;T}EufLT@ zq_$Yl*`a#zn0h`4@E8~>GKmq;X4zWN=bqhLM@#=U!op zOGto=GWwY}ZrphPUdq(QY#nVJ`TRNhJaepc3ff-JmEyefP5{Ip&?~iHyEw}5=Pph& z z!OpVooz3VW`s*wF{^#g`z$XiN zo13d^%X6Ku$Vh9R5L8Ib_)>SzeE#;W3xymmW*XnOsI8<2|M*2nxsgfaB4pMX#A^=a zlTK(vh$I{<8U2TmdiLd(C@74Yf3~GG)7GL zXwgW+E5kmL?IKk*p^S$r=i=(GA17np@bgA!0rMI4%hcQ)%=xm-nqdI`NF_d5B^gW! zT&TbCNHQ{`Q$Nq4Gp+_)q6@zxt^M5h^T{d-+wLNWG{pJ1TK@Wc=KViEUY(@~ybO;= zFYopVgoQlLu&dsCJ0(R+Rm$>Vy$OT@)7Bd|PR^tq8Lz4HI9HXakWpE^)LU12D5p-> zS;j-TuSegYaDz|H=;9H-Uk7NFhj`BADJV@RWl5M>qn!x_Tq=*02c2n-l5&NTY`GY3 zLqn>+mF|-W%+U#7E-i9=ti{Jt7wHqbSH3FCm_-rRO7d0yOox*Mq@9+92I^z=pg&ysR&E1bwXFs!pOx!lFfj(rBb?XO`dykY^+|P zn&X$Mz-ieaZ~3(jy zqw@A}dQ8TC8gA^35aW7cT<0#Uu*$&@Cshv7S8qduyR72ciGBM@jjpU79EQbQDU(Q{ zl#R_Q^)(Se^&ZO7OjgdE(UmcyrC*~YidgWfUxghDkyENPPK%0qWyq-4DOXVaEJ`J3 zf38s{OV#_1@(&;6oQ#7bA|TDV*hHk4dkM77&!P(gkqBR)!9AR{xW}ya5?~@Wx3Kyx zNI8yz>lO45Cnwk*2A;YDD0?vppjQdyRhqVr$`H64FkA7w2UcaOUa@myvMR{}u+LPo zsy3>!T#I(^jpVWHC^ss)^4Y$A&|K|yP*U@X>~$}+XknKn)kPFTZynbd&+(MUXt*h< z)VL~u!Jm;MPm$`26ZZF!y!Oosqm2Z4>)Wx~{%|OfKoi$)V?# zeW@u?u=>q~7}CRiS*DBAX!R9K$1C>2x8n0g@`s!kHQwIg(Atw>JIqFzZ)^I&n`$eY zI*Fv3qVQbv`jszl8yWh)#ml$ak7d@XKJG5`F|3Pt^R3$jBGX>j)OW{7PTcjRmNx8P z88Uh7$<3`uxjW)0<0_Z5wBkd0kI657iXMwB10VK1?w=Y|#yFgktB24xP;y^UJ*wJY zMJ2nsoN%OE`%^T8}rM^oRuJyc)L8EVe^5tL#cx7Wtyn5Q@uR9 za%~Xe_p1?*1A&Kf-5-N+`GOjA1Teolfj}Lg^=J!v786s3KdCsiw$>Mj9Rt<578~ff zz}uhdHKD2#Y=4W)0yE$Q@)F_5>NZ))L{Hs}%0EyyZS9W%pm{+z3>to4RX#Vv}#MnG@^%_jj?O1etFR>4?6_Rva4VK zGimqi*gJct&pI*#dNjjHmO()rw^r@?QkC)(o6{ehPQ#B`3F{31xN1MkLnGsz`}%xM zSxmq@My;u)neG1(XFT#_XQQrrjohzi~=*TjZ zb~G_*dh7d6z|y-TlfijuHkeLC-+uf-XQANdS2ypJc~bfvklsQr>*Xco^+0sjeit6j zaYOxQ0rdV)6o8P7qxyzZch3RS^NZL`?Bb-mG*zjR0^AB(elAaty=wXS-lE~asUpG~ zyv^=uF1#oI8Fp`Y?7`I$#rL~MG6UBf+>hNjvq&gW9=76d7%lVH=FNhVdQiS;HQY63K`4UqiL9O@hDKYDK7ZRJUlE)#8F^0)fr<~45|UD`}~b0eIWN6Plo;tkGD}7I<3X{eS?$Z0lKs! z5h_KAZ;BI|y@f_?Ih~8W>YJO(YSO&D3SB-cUUwBe;f85w$RKJeUX-R=r`@RAC$rRcXBe7fY}ee)?d671aV=$9fFjv%M;HQEf(^^d?GKRXQN{zrMQVui zau6er1|vOSFusU>6s%dm*7gFa#|471>l(Uh!Oed!o3 z-s|emP}3p1*m3wF*YV>ls;s_hqp9^mrOU%g7t3{=v$d9A7R!34C>S@5`uK)rWqJC@ zQr(VznOERU;~N+2T|cB9Co?wEJo?vP0RaVn{6S&XWTKU#))>lP{p3?eb|=3@8~3pn z54M_i3B^kL3(B^}%NGio&%SHSa+*@q&RvykNU$5N>&V{!F{x2IV3kh9y(K4Y%w<8l zF}eAjK;)M&CZV?3&KX+MuC`g96>mPXqdWYgc<9x9&0Aic5RjGO@bIwSxzi!$N8pcR zG^=P0u2o#QPh(bo>ZtFzyJ9~31GWo49s6$dRQ>rsN%I7;((8WZ zeZRGYCy@cuot6I9$_nA{73E&K9q};E{LEltC15Oflz;a@tECEeql0jE zFqpIIJd)UUu|4C<^gW76gRH#emsZ`C{N``=a^;-pu{=W0?OmM{!fW}Cel!)Qe|V@q z=SZ+!37u%0{d87sK;22on1NXLl}+};yEU?|sK(UST>4ww{PVQt?PyRvaAkuX&V|a- z9YRg?z!0*JLL`r7=}EI+4Wo1dE@1}uQtbHg1;l6IZ8MXTnYM|22{K7dNzGQZwau=s zXKvk=oF2-+g?fa1Z+Q31;n&gPrd(n9%jIXCHqLB1eywlf{f{qq9yze)crNJCc4myq z`v(OxOVrDX*x%vOT^{=wy&__l{AFTd{99_P8M)%dD3yYO0`#u+0Rtm#C;v*zTNFTuoYc=8fI_fP9%CWlfpjJsq+Gumw$Ar3`byr`u9`NI+PpVHPbUdb z5?UmLDGSC-rDJ+nL`3bd$FHv-Oniki2$;OQQ>flRy#g$Pk%UR-O`6RVV=}}W<9m1P zv$3isC4Bz8mZD%uPKMv=;j9q;cZ`*@+=mc=yZCnnm&I@Ere8>*`kJ-T(aC9MyxRqR z1+?6DZx9MxhdW;B-6XDO<|9W4mLK|Z0gSO`O(R;WDYNFi?v2yKetrjQQ7BNG?n?{B-mxp7T9mTA!7WKn-#lrRk9EA^EWJh4Dh_UAvloe&mDU z`5%M3b`_ScV7`n5f~XPpK{x>* z%Z*9c*iRtPAWFB6+Ge}DlKkkFjiLG;c!lu9e!iv~ruG_1w18G%HVL_6YmxqEGGQQ^TKiaun>`yB_8S5P&zJ*TZ z17DqE`ggJ!r`9yG_ziTX59CcSm%SpAU`jBjVT zChy>l76c>jMbM6Sn5deaZKzNN`~fBeyc+s{%h1mY5L>1Mm?{_viGXea&R$5<;$~C3 zdQ0!j7AIqU8JU7?YdsEndhk2|vMD8@=DZ33d4cN}FdIPN8v#d1rb%Ym9Z~oSJu8}y zd1JkO#A6g!%F1-`7~IQJVA2tV@Pz01uhsMgP>R+!*_b%c}l@Asl zgay3t<1BmZmRHJx*e`31>}L|;93OgiPPM!1Vs^)_6E5)*=h~l7rOVhD_RooX79=-g zh!X1{M~J-amW3lMQgU({`3>h%>}+h5SOdW_B*;B$`s+{g6E`{K^aB6`Fu(oVH~Gpm z9~YYMi5AcCa9v|ldK4V&{H<+8+%w|dU89P(dzVLSv7sU zW2W;yWjsL6`FU(tkIycy2X}EBydqnw3M|~X*@xP?@Xtg2-LhJ?)5f4@r>aN`H%?%a zAz$q;-M{`-&a)zmVCKS`5YaE+a`(*TkE4oDyN(&j)a~VXJ2PJF zP#dn7W97DY*W;O1zK8h5MR)EzXlr{pG&DUlv|L`U-@2=M?R}QdJKgr|6zvqXzysSB zT%2?|N{T)@l;K_R4Y4OZ0XypBBvmSak+;N*mEXaZ)4{f zot5nE$mPO^aPVG9Js{@362KDV@a>VM7`D*t>WZMkHza2o0 zR@_M2%i-H)o{i6*3CM?i9UZr78Wj<+Dnk7{rJd_Nw~9svb`4iSqM7cQHx-thC3HN7 zogYH_RNJ3ls%(7m;=bFGyGryKOOLh4mb5S1>u?Ud9!oy(_FZ0qgeOqO2%HK+HSq6; zg-sDsy$%tfrs>&L)3J`yGzD27?+zX}XFD#US#_SdoILS-(i7bn)6T^DzWNwN>$dNr zqD3#PUOHXUB=*<7{Y(8r^*#WgL%6@-ekeX63S@e^cWqVu^QU96vecu>9eZn;A!R+4 zpg@tV9w}mLG9EUY=jV4I{PrG8mR4upl!<#v%zZNjgQ~yZ(AYX% zg_-uRzotz)XO|uy@ngTptlNX9eZp8>a!o1b8R_A{m9>@LwT+I>x@`r`_8Vo9BPx#d zuHD`#Hf`hJr;98=89*(zGVKpy zevH5J`KQ7TvtwB+VM}^@{PPvTi^cdfjt#D^(*%3W)&_U!;nvshOu`{Fe_TyEK|$`4 z_AKKc7HNywhKF>85S-ek>&_&rUVjOd^wk#*S3jEaZAxo7moiqLIZ%9iGsn(KZ{Kci zm)Xtq?n9}~3QNmVMS`Imth)WkT=uJ^>n7P4S*y3E@fv>(%bhwb>iQ$`0;lj(|7VZp zhd&K5>&h8z-Yn_jV(lFaa(@EB zRX$LA-&IPwkm8bKI>pESVb*%w!*tt|ClMGt%AqYVoK;8Ll!$%|T+-Ci?5|!;RLRii zSFAJLuo;;GH$Q9gG;d$%hrt|q*hKctjz#-3+L?8zq}RJ*pBlNiWXlXZYi*sHqgKqg zTJH5D(!IcXyz5(@P_+TCzP^A$*lf?zgiC&Y+Ec!yX(m<{uEU2(zI598N(?`UWT(eT zv*jxjqb`&Wcwv2R$63lF)J9%vso#FBnLugY(cbAXkC9b7cKDh%1@qbtCC!Y&%X96$ zopv1bMb>nF+GCjC}Kn~x52+7 z+_q|!yKkZzeTSRbvYJNCSMMSZfBpD~f%BYNib}Sd8v^ELN2g&vS`1sRvbX0?R@33O z>Kb!fn%up+eD%XGB)U6t z(_gUcIqsySwy(tN!Np@X{hR4sKf?bK;8;$lhfz7XNXZrX6CNHvid4J0S`@+<kdYZpKr8dH<#3xeEcEr|fWih-@_cIQ=(-8@k zx#=e6Kz(wVxQ27-JCVWYEzOEtsVgbA*86O9UN-o;=A(e=>Xx*$k>=|*w%`7yF>&K) z&Xg1#2|Fats{WjP*03;Mo}T`M*@|uc)~(!Z>IH@Q?IOZ0lQYG6ZZ0<$Zr#nDoAA>p za-4++#(uuU;W10wYrD95G-DPrWrI*n1@Jt5IqkLiB>7G~tc1%%Dc5Xg)N>uhRBh^P z*M*i4UO~B#%xfb}W4*4|Mfp@l%#{(G{OGL*KV51XE-wy9x)TuTb^6oCA}1}?Xn<^R zN3Zk{Uw>aeJ3r?+`v(W<(MCl!Uzqn{aZ+V`DR zyVx~Mv0q|tWWb1zwWd{h3|5nU$*P*?66!*qJmIuxS8~iy5)rw(AX?B|T+HOQ>a;x4 zJX~KNzgt*Emd8MUwl?I;2o4g%50PC^gE9n1P4tt4>hrFaZyEJ( zT&-bYR}K@%&dr@@N^S4nusQZF(g{>kpsn>X2x+9Lv4w|+Hb1A!H@rWy;9em-C}3@+ z5hH;J1kQ4rOfxJ2uWRQRYU)sB2Z##WV6x0fDps}Wnd$A?MR~DXQyD63m*;Pbd*W@4 zGDm2ot;VXRypw$RzD3&CTZ{i%pJX6wzDai0W(wAvUMiPW&9mhs&8&;zPb(CzI$YD_ zzkKWE3bHE1Wn?PePg7WX+^x=>QE)ACB-QXgXeN@M(B2KkcR`$Gp zEuFN-kpL(`gBQ5`?yLV#GNv>Tf(f#(YJP_d-G^hKYs+CiUGNU}E(4+Rqr~ zqPOqIUT-(^JC90*vDO@;EmXJf!l7|1H{W_u6!n29ii(g;LRNfod|d4Rqw6cfqTHgk z4QvDzDG^XWkd|(cuA#dnhVE`uM5Lt|QktP_=#Y}`?(VLk;oEr5dCz;jo^O7fAIIUD z{p?t4?R%|t-~M<ov0dwpD@^b^a&0@8eD+rBBzG%jRqEZV?97mgJ z?r}^pDM|H-0%$u>z$$RCd~#-amaS59wqv9`GH`jpb3EfoahHE(Sc{C1BZ2@w?kWBS z5z(Bea+Kq4DRPUe)M+0i41OK0FVAgkfcCS@;h@#QcD?w))A=z3K!PV2-A;iwg{pOl ziDVwsJF`)t~G~I*k_Ma15?0*aKG5?kq%~vVV;~QT2^s%w6~Hdkw-OX z`lZUvQQ$=;9Hco~zF$9-(_LS0>VW@xfmXs?XRi3>ZSG9v0#GwYM@RR6>>)O|S_8h1 z?6oaTR?I>Gf_ex`#+Fr?DD68RZs5P98j&MC39PNzJ5)e|A(npYhZUo~|dZp^A zl{tE#9jbK{Qs9*biW4tizXtnyz?ji(ywpRR+aZ@qCrU!sPn*Ewz2=GKCJMI5qw`+X#mh+xPlP?HU6?3+p;2Z-xt5Z5LRa>j(hjq1mab!m^ z80~|R!}DqvxsA`s`if-X>T>!jOT9oR=I&!|lj*6}I^I-S++tOOfV6z2Aw%xb( zcmP|k%hg35U+Tk51(W*p1Okn+9k5~S0;mZ>2rz7{H^)5pm#3K{hjfZ+&8NAO47y*` z)R+s+cYL}Xs-;*3P6yeXNYL4=w^%r(-}Eg2uLNH?|7Bs`xBfmLtf0pg#L8E@^oPT7 z9C*XIX4@;rwI*p}*`Wzhj_&Odxpj4b+dr3v=K7X4>6)ZU?nQ>~r!%J=LHgnY5{&5)*G z!9){xj*wFv27#Na#!Vspcr1!cFXB4D6;trJF(OnhsDYPoVb(I9>`!hAagKqmeaewT-7% zMP2KIUM#c%q;?LNW0n=$(6hsKYA17ZMbnKsr~L!x>MTiYcB>@$G+$~h5MizRe++4ZSqi#TY3!J|e@C>{gPww5@Eq8Hpe7^NQMoHxU$vV;3#u%3)=Zf-j z98`c!M*g3LDFdCQpd!U>Vp1uP@CrQwj_t5vsGOw;z2^6*zNryzbTogP^EI=$(A1cu6wQM7^LUL?rz3_7Z|w@n#Vv%2_sJ9DOWTTrMfv{FxA}n_ z)FI$q8t*g%+2K46rhO?t`83ejSJ#yrSbGHu@_ni01OZmXnJASI1cJ)J^<`8%#SLTN zu@9o6_CcTWOQuY>4?jvPEBk;WoT}O3OG$9Gw%OJ+ayJIx!>LjK%a`~~dmwbJ+H#?d zjbJy4=yMGp$|N8-tG%ijq5d{l(h&>HBiNTkn$tbgwP5dQczOyW*ANz6u^1*p>cDLG z_9m0}_KxoEo(Phah%FXmflN^U(!V4$&Al8Yaf zRDFYo3tK^%_HUdH>Uq~B4MMkaIEOp&`0Fo(@#=x|?#&aNfIi%h4_pio(y~!n67y8_xJK1%foA~&4s%Eua zakQqBrzBKVT%>9FEaAByRQLAv1%xf)n8e|-xa#UPxaW`g+etgx+B4x2Vv^<=i~5o)!$+hNCt6(;`DlYyODg4d`NYPfs{yp8gRt6 zR`&7SmOo4b)1L3{nQu??6nI?u+S~Vy>PRd-^3!i_25%NApY1E5Fz(!V3VXcrjzAL z&kmrimCu+lde>N`4~=<&iOKF%m|PBtW3f&tmgem3-`UyrMpt5IXD5FSE}_Ce!)Tb+ zzDdet_;R`W=)m~!JBV07 zfyVv$vo+;5p?v%RyH_kJWm!8^GwArPc|M;k>}|h!+AB=xQ(dY%l&v!@15R|N4+fxl zi$#oGZjSq*zomUW+ZfF>$jKSoQKYf?5?=%LYu?#Wx38agkdF$3U8bj}r)JjVVJFL* zO;NrUt>BK4k=Gs91z_Og=yr{AYl=XPF+H`rX6l@-%DLy+`9WkX76L2@-7rG-~$CKX0z?#g?>6iLqB3$rXWa zN5$L?bH=TU%QFBOJQ519mVfI{LxTbQCi#g#SxIX^P!O=gFQ%rDXlTxH<_Ba6dXcWF z-5h5}9%B`DUJo%{{eLBZHi!~uD)KpMi+DD-#Zli$88!XQnE&&B?hM4*b8rJCEcH%VtY z`1lwr2zP#F1%Oa!TVo2wmB|s*7*Zf^PO-FphPTmQFc0Euk5HP&crVtH`d1H6%dch* z4EfriAh+sLYWxQ69FUz|*8*RE{x&)wgYn|shmVzh?fy4!TR(eddNpu)MDm7y@5*4C zdm>r{o95%y?dj@D->@)I(rW2c2NW6kCmdKCn3Osel*-bkuS6aoX%ri$p7Orms7^{E znxx5ie@^OyM|a!zTdsqR`w><9*H%4M((X{3X_x|o!NU{RkoL`q!o|>wGafyCMLoSL zu(^+0f@44o0b|JvN+1*siGJ-5Y7s0<&9%rlfLISl$O)5)#ObIT$)hvJ(HTBDOj1WnM-=VD5GE$n&zZ_=T^vdkbAZluO#9;UaNbp0 z*g;iy3v2NA52Vj_MGK1}hbI>nBH?gYmZF`>WOSX*QpT_OAJT)3dgS5-esQPGz) zkzyb7w)Y3*JLjRzfr9?Hx_WwgZ$h~yio8{$jsOL&6eO%&?tZ_mcZpJ zw5K?YoieD)U0R4@3E(Y0% zfpHJeW&#R?DiWUj-Ur}6()Ykgym1QB+a7>yXi3)S_}kOk8i>N>q&iq1^^szDOb~kA z>~UEIfw)MEa<#J-#9!z108H||M~|5OQY*4)VyH2FPLEGCCh`;6InG84bY$}+a5TC!v`Ktjg;gB03qTzolmba)JicK z3`&j;QZ!Y6F3V*>I&Id^PuNST1{b<=HM&YlenHY=Jo&-BC%Sxc3Q)iV1wHVb9jS7f zP2?Y_mj^ny0z$q5L}=a(r1y=hOF9$s@bScQ4o?E`xHdNKGmlRWAtkNp4`=r0o3Z_I z8}tf9LO>brCd_VKGbOEt&2G(3Rdp3mKbM*%V-<#HWziwbPR`xq*sEog$}Rf(hCog7 zixuug~E%_tbkW z08IgPs6hzP13dQPq%F<-Ub*ma1^a`QL29)^r9AH7@0L{D?iTU`$L<%bB;*VGan#h* zi`Ji{4Ld$CGBD~DjNT%UBE@pyVCy90IvwhYi+(g-)86tUennrG{~~fsL)D zG|Vh1xy8(^Xw7zLOAus05srp_zYI{WK=nxvlcbQ_3G|j9at#X{)=L8NKeO`Y4P&7I)CbQn6R|L2{Tv%D9Yh5bK)Tb*D zLrzPZqw~{AUWviMsgQ{1C}p6nvl+*sOC~Y*6s2_$ZFNY3A?W~WY+N5kj(XSg)qUeL zlzVr%^Q=BWJKs(CT%5|d)Tz@Oggcz!sX z{o_?F7BwRWA89bSrg^Dyc5+Eir=*%iamA#X=bfFL)t%^dh1daWr zXP-WH&+i_)(uf-8s)=`gau#S=R*Qi>UL9l7T@;gebT)giy1``PIyrH?Mh;ifyTX4w z$Zs4i%V$J-DIi)#v+X2C?xbkt@xexyA+*W|3v=P zepjx6^KFoZP7_9liwNYvL|z$gE|t6mjrKzks(HPWmBV&26Q6;Dw_VrQ(cv@u6d@_e z&u+J8Wu*S72t=+ zD=|~!V2cwG=^_89fX%4*Hd{hsQD^rQ5R%4oWfkFEpQqnFC%+fO5a+i^Oaf%?jEn;i zintS#_JMkS?+hiqtnCZYkd$xTkv>gW`ud7kSoRSS=0ihf8XDfVtA(s2dc*Zzkt`PO zKx}KRJ|7WKpWW6)1N6&HHplF=wQYbr5I)WbLLx?yh6dnz8h}a}xm9gj|An89P7st} zY;dpw)pe=%5HiBJbPiDjqPkMPk^T0kX=gTVkg7?1=6G>VFn z16W1H2Zg2E(}@kyFotgP5oy5tUsO0?G6`zN2KE*$`ul7(`QSAPCL1ZDZWrhDS2enn zpI_C!0tnCF=i|j)a7Fr9^A^xme*5+=R74;3Wd_AhLoLKr3NXfm;vy>}nc*?jZc;q8 zKA-{>6>aJzfjd3jZv9ec5OUMSZEvB{A5}pmDpYRTo|-Hj_^6B3{>9Z*3#qsyzKgxr z4Ey~M!+uI2wWr`-<#AOWfG3r!Y-N42Cr7EIpQq*kT_37!Z3X7KUA<_jrx03+MZ*cY zc-ME0ruOO8xIMdb+G6X|q?#z3}|$5MPObHmiwM0%Z%T3Jd< z3xl-337m($iil6D2ydXjaR2_3hbXtWwTPaN0cPU9bace$?b_@Pcq#>GZXOE6-(DX% z3x~lRve?LDzf6(Odxhw) zCm-&Y)P$@zL+#`c32ls8>dH&Fyb@ZQZRMZgIMBfnE~xmV4!1Y-m8Sq`@f zddhXVU`63+Bu4V;m2tqSYUM!T*uYh9>un@&sbuQe9VsA99DFuHU_Q%emm+P?@3^GM;KVs~}?_5HQP zW7O=9&Qcm>>27F7S66Nkzu*&r+TXp4zcYYZG(LCKB_^gUSEsXEYDZ#XGd?E`xK@2Rd*0! z5RB1h#!}!`-2ksA_xVqJ<)+;iP;Pf?8g|1gVli$w-my1IfyG9Wa2D3f*>ArsHGe=t)AZZ3=r%q8F3G5gaL z^r5|E{9qhYfBrpGWpr&snIK_%ehKNgt$C`n*}@YKv@hCYSX{cRC4Iw~7LkGjpC1AqpIo$Ng_50}o&39pk2U}2d-?e-%H_|X6rW#xyG#M( zZYc&*|IYjLqu=JYCjxG-uiuRE47g13s}g)(u)0cVpmdz?!>fO-%=UFL)q5?&c6l|1 zqGT&TM+x@%0oflf0iM>mxK}MB^Y<6V=$S}JP{0rU^?84BJ0cW5$AuK&T2px67k`?rz=H%uM zU3~fb@i#uQUoTY;cAl8{u#nI*5+m}+ZS`V6-1PShxan+9q40tN(%)TjzyIo=BQFPQ zn^%dBHrL9(@z2+Zb#)TBP5;-7J@{wFu+q|i`Qu(%=&3am2K)Ik{J++TiG@i64@*(v zEc5rQmR|e2zvrRhuW?4xhc_(8z@kYZF^b{gg8xkapI>BZA2^1~$)pX%$5w<~1jKe` z3e(rpF0aA*Yh6~dkaB#}GiwX3Z;CIJmAZQPIUJvGecS7R0RcsH7;KxKzBo}by=}wv zjoI8Z#ec0?nzhrw@h*tRR;Np@Hr{sYDH#|jLG|5P-A>(vevQkeYsPV&5HTI5FR_JV zklBVhpP#hqdMvN(hYHK{^apbD*d05aCU&>7W^~k(eL()ouus= zxPO>n;NTW)Q>m#wayeJeJgGi`oZE{SD?Di$44}<5cXTqFU&o7CYxy=jpBO*Byi7x%j_gE`IkbDbh_b^wCRN3^%(AMAJmWJ#Km7JqcAEQ? zWNrc6wN(+5Z&{ME8!8{ir4Z~i=d?DPlzeXwxmjLhvt0fdPZrbH zv$q#-xkH87sIRJ5cPW4%r}O+**Mv-R{pZiN3_6A*Ep= zeo|+Bbyo^eM4v>sANOj!ZgvRDav*&v1(Q#R6@kUNe-P^`z%^Kh%`?!%J*;$hMk1US zoR>Bx{PwrCe2?iLZD}bht6|AFSz7I4m*+3kd5KKMraw#VY}l6xq7mY$d9r_ePCpsz z=i*u;)YZhpo#hz25ACAr?p9H8G#^}!RgrVlpIp| z$yuH(WRI1$MFEkxBq9rW6tq3mS-|BbeB=&u^Abj`FYSC9v-uY9>E(@PBnMDh09Swr z=cto&taQ6>s(Ko9h#kqv+J?l0SQ`xIny8fyg`}kH<1xo-)Yj-*59Nrf70fxhx^DXT zaa5R&BJVsZ>4-@em33sb`|O**&Sc&eR@Ry+AQdYT2l;&a_Vu0o&9lV!ubw?@A_|tL zDr90XnV+RAaoA?k)>bMpO;o%xQytq@gd4XH0Pl1_s#%8KK9|_OA#UySgG3;kFsDxfyS#frfDk0UCU1qDdn;VW{5VXydsde^|IescI z&>Eht!q5On)o>WifJn#Bmd2KReSg%)(QaZ4{0+4p7tcNQB@jkq9Y5f8%Z!PA2Y%;p zAKoIK#{rj&vZ<}pC;dIg-TG@dI9Hko_Y5 z>hHbiDWeHI8W@LJ`D=iMN!gTHi3WR+K|f7L%|wpOrTw(AsjjayT;q&+Iv8zpgX4_N z{xsfp*I{F{i9A2chI)1Bpz{FDa%;Mu$?(0Oh29TIPGi3_?8NxNY08Z0$kp`=sOp`= z-3M!vOYOyv*9Ik7xGp{e8bZ<^{L-a?#kLRE`cm~zk5BKGn4NVV?Ch>Ctd6$!i^l{e z%&9Qob8?kpL&F9=!I11IA zlBj~QUCoqO*%;WakJSyz)670Y@vLK{)5gJzk*@ttLLTW8hfh6@D8zlT!p2tS$&4X{ zAe>0IQ@XO@?O4X>V4r{b+=Wt@L;UOo(3{Ig z&>@O0p5vLwJ+fI{juKt8E>I@v5h9usZ%kp`)sZR;y;JKoW@C4k<+fuyvbB(m6J<<)u0@#D<<9Jp zVD_@|)I~b|hdI6~>ejfbo&Ae(?O$IcF@#FDR}7@$DzmAeM6ZL`S>cfcBk9nI`YGQx z7fH)#E0g73aj9%U&GUt$T%e~64^shuQbEwOU4*K9Q-F0AMWzg#y_MkLzEHmAo;ajJF95Fc^t7>VM_5VD`-M`7*(P-fLl8cr5^O`8!_-KE_7@p_XlpgbkJ)~-w=N|6@`=1dABJ)yw_7_ z>F9#PEou1<&13W)5(>@_yvoBy^X!lZ$YH&?6=cKR;??1PwU`>NY8P9x+2O9%(wt;c zGq;x$8adP*klJr>c)7YNgFbx$9FluS(OGmRvwANL@({@q{wRitLA&6=Ak(`CXhFi!Z33DEL;Pn>Ilj2A zTqek@<^qK`K5c!CLA9K17HJzcnc{YGhfoS}+x@o-)>2{VlCnkL~mGm{FUI&eYDwcDX6unwRw0QkQl_(E7y9*WM-z@B2nN z+dUO&@KZuD#qZIm^;lS24j1so8H+a;@Yd*hBeI!vzamiAw&}`()HHAh#t|0i%bFT- zYoP=hu@$F)dU^V%?Vin^#qzOB8b3$6WKG4R^Fv#s?TKnRE|xf(gd5nWzlQZhv`S|Z zH4og2;{vT{M%VsPy8Bai4WFR-^C?| zv76h_4sga36KJTs$6xj!+>sgBy+RsQAtiXLg>9#pn5N~JrmYT2^*ULyw_H^4Pep$e z@O8$r2XR#r#?@HtHTKrTZw>J!0v(LT?q>JyrGlVtQq8GyW4-rECgriMAAk>ORYt5$ z@=T@!?+kiP{V1`j^Z z?3T$sm>z6(?Z{`2iM*%L9(%DSMF5Eix1zqRG!$ET~5{|Q~de% z!@2hT`L(w3^N~m=y6!FEVKV4ZB-V5QMFqxa{__dVy@R2oX(m5-Q!?%Tur*5|qXW4c zXmquq4 zCKF4~C~w#NE@ZL(0kK;*zVBNx@u^bAzNued`JYh7#YxdT0}ds}oBW14?GUPJIk@n&z}})At|z79J=2ZVR}-6_b7}mRyQTv67ah8Gis3P^;HS7avT*1OUFHD`Z= zkq`Iw4J6~*h(&J(4cEGz6#|TFy2OJ6yDy7?M58?#V`MV4iRjC%zEVU3l8Kly{+An7 ze+~w?M;|b6KblA1j_sLf^Z$TR2ipCmL|c`sRf&SXtDA2uw1TUwZFer1cKYyEJ!cn} zYZ=!+0k-R0SpvOR`{JHxq0)GXmp6Jpcw^8O7(N{MT+=U}E2lE}f!f=dyY@C+Y|Qw# zWdWA40;f_09Bx_i1YHmPV(7JYn3GfapUVmq zpaE5foM?!tBg~n5{vv0S>`2ny`rVuIiI0*r*s@odG%`JuC;+ipNOpEztJ|EX>iB}< zcb|}s2rxaZ9gKI3I6}JPyQN}#xSB6J2B_qKghSS!_d!s7oh}iRMlR<9avwb(=?~{x z=i1we`EHM(^!D{5M{*yCsgGfUU&$ix4)RkQ2wJdSAhJ63#e_e+6I1tO*O6RCreb*{ z0r%hI@izeXSB5bh%v#yR$DtfTLe}?JB^M^Gt(TU&EvpwEY#h<8dUL?|*V6i@58nlc z;+v%AViV+{dH09>iB0@I#he$@m8B0n=~Bca+p?)988)T-09a*h?R>JA`X4g}vKk^H z_aocE6rpKq9s$e`Y!YfVm8MmWC2xC!cTr|lIer$jC*QQ;;KvR{?0n|*ksC-eGt>Z_ zdBBIBD=l*8*6X0inAmeVGT|e=Gn*u8IUuM6gf(pc07nHiHKzbkFY8OR$o9jlOL_$R z@qv?T3QXQ7DN}F?#n=g!rp7N%9-!RQ)%ohMKGfJ_Iojw`Yieoi z{&Q{I^C3B`aU~4;B?SeEISAM0ir_Le06n!hEFOa9<_0$vQ{+c9 z{JHe&JQzL~mv|sSYJ!D_P+Nb0v+Z`59z9T^2ITPnf$V}lUXK(AHHW?JZW$c_$5j!< zwNc!DQy2hue5EC7S1`w?nh_W9MJmKb*gv0+p@@;GI^T-;!u5<{kAga_bfaCc+FP)Qz znW)%>fnP{mSXfF}nAGzrkJDark-=*Egi!Z&Il|@`9S$lC0;FO2C#C)_O}ud<<2S96 z8z4jkqOYe`17^l}Qr%hOCBPKrQMPStfEtUOi3u;h%Y_*<%oD4FkQ>)Hq%&}Rky4#= z`zNx&Wk6HssV5!i5Y|#k%M?y))1&wQyXpu${x95~h=oN-Ev>}7cg|Y&*~K`rLKbL` zAe<37aF{wxpvFjgg+S?RfCBTbg$@TrdP%P7)qn62B*NlVZeJXD5aB*#b^s7g(1{H4 zoqrB~ml(R{-O*9J*iQY1zFfV+c=Sh5w_*?8t?NS_1E-J_8E1hB?~R&%W+#mCYkDIb z)Yn}RKwt)=!14oPR=Iy}7(n-M28bm+p_wrC3L5RV*Kp|fyUQmRE4ldU0RbulC(6H; zr_Z`Llv$_%0tiYwSO3HVK#h$9@HTbOHTYtC4v95Gxv8$slfNIuFTlfXk@w&#eH;6+uh6g)&SPOkA~nP7d4)jO?`)&R>s9&s4TrT4D52 z{y`zx&sN* z6Rh#Lz={CEfd5z#K=)QzQDNLRpg66S22yy8ZaW{C>tY#y)e`XpPW?4Vh@I%%VaMpz zkw+(0)0`5SRG(D&@4fQ;VQ5ImpZ2RMC=3qHNoA$|q4{+D%##Zj=W16R{=dhURv}=; zWrOs%%!6kF0r1oWJ3KD1Ca<2D0`aOp9VA#2RnQO)BquD$$U+tlRyNcfO&(O^{azE; z+OwuX|DLP!H+I@rn+m6p&T3Zi_%;$ximLQ`(7{JW`n=Nvs zz4n-=bHZ%{pjOO8zjF=_G)8Q8+g5IhgqF84iT3ARD9n9#<0g=|X5 zUlMk$vHa(x1;@ongy=Q2T~{;cqu~QQ`N8M!q9{)pZf6=z-uV3LzHGOA_*;5(sf%up z%jcmB;VC4iDy3F(MUcf^N6^+j3Aos-LaQ3yh2Tzf;Bt%myfcxwf10ZfBoe4rZ)JB?BW~((vTF zw4;0dP7%F`Y9ngelF+@}=J~}*u3kI9k>JkKyZ(ERai2&>Fxgb2eF))N(+l6kiQ1z&|1B zo>bAawApQ@pLD(f7J31|JWlS}g#x@+7)fMKN;o_&&QiCDu_7B_EY=Hov~pRQnYKJ+ z|7YlYK-V3pnFlq8`ZQU?4_w*@06DWGN2T(3OM`)}h52i{lE_8PsH*~CW?#R5@dVnM zE|2v_enqP}P2TLLphCf%8}E5ca{(>kuc|gbrz6Hd=P%s5wXZOx<|PSZE6g^(ea7<( zE3PGS4%th$ayu*5ExeEB&0unD3AHcT@+p=1UwO{D-KX2R?M-zcUd3Q>9>82mPZ)gT zyQU)um8tixKL$`s=JQ66w~*sI2%19*!BS!y+AmsMTpr^x34`;5War1T8 z4gX3DDD|>1F+pd+!=qe94mRMqjZ=tYFxeRix=+Liw~K8Han)43+up19R>GR#;<-{? z+VC@CC_=5!SX;Jm*|zcpOdX7kYh8Sbx=w--S!Zkuq)^8)Habac+iO+FJmk$c(ord)ir z-%wT~$#M79*D(pHT-wZI#bw{~?nN_q(+?jpcrDAWcm7`=cZ=VrfDZ~ywb#pp5E&`M z*5>pOZkyjEH{+-~^bDjr3=it1JY-eEzPo5&p2H?nb%6Z&mrg~Ewb0I`d}Wj%VKl;l ziX=LU4yobC?k|q$Y14L)A!=+k-7jO}3!xi(E4&y&Cub+Dj-~mB=GwYYat3CN{zerVN;&$%3ahts?K?J6LVF3?)Dv{v;eJ{+h z7PV~6%5WF$ya{Ryt{I1kKhD_xdHWlPYXkEGn}U)Ohy{s-qBbq_;gAlkr8G+hi0@a`8z$dKp20A1zY$F@*Ul4FjWv5RDvO-$lkHvs&=AWg-;XF=b8H`Ri8e^ zLyy*~Zfp$i*SxyiD&JwU@^raZ#y%4Lf>D>z$k^~hOVU!uyw#msAyOB9?SuX(m^*As z`nQy@*8|;YZr`qMMm1w#&M<4SRQT*UCW*fuSz*AJT)L#th^2@^;m5*O$%+oCQl5*quKkvz|#4-i`q9IWe4 zQofJjFX{IslZO%j(VREBic7ObFz+;d!Z>$w>J4Cp2Ah0i43-?rV{Ale5U1*BbZjF> zqZc0hasICIgb%z*Ba079r139Mo=#Ms(bC8PvdMGgrrLkZEzr3)8K&ukx-ZJE0=IYJ z;^HR|!WAQ|8>@bMkDXS(S|>xfh3tM{Ksckb`%G39PHKfZd09EDt+f)C_#pn6GE5&o ze(K};x*|5^R%{$zPj7EaRI5$_Tm_!{y5kL0HqkFM+(WIhW>TZV!P><9b>bxj)`NNj z+eFQy(dyIz9r=ak8p>}TTFS(C9w_i=(?%}RJWhD}qoSD3yJ69K)3o#RT_?g(Ce4$@ zQ$@@-((CljhLkjs{7=3lSDysO&AX4vq>qvxG~;1`=)N-<>J&+A9uF&WXji5Yw!0y1 zvc2$OYfFb5>md#)>Q_MGPA5+Cq>*NAe|L0m@0``Rjd?2YS+KMMbEUFI`$K?E4u2Wj z*nTuO%^65tj&XAED`fd7Qr5+r8&|el+Iy%U&ri+8okFhp-@MFmz`tf)8M;Oq2A{?U zL+Kjbp}Fi55`w@4aM<${(RC^I=3g)~!c^-Sn)O>qh>V%;xS9)Voy}L0)wow>50xu~zcr78 zjSo)4n)E`|CG&D=1dC(RDF!jE6|DGjY<_r0*8uMcrg6X*c)k&D}z=?*pp;6j)^ zQ;oy!+>)ij5^3VATnZ7iadvktnM*cSBPUSzRmR|OoCWweb5K$QjC((5?e@p9Tk%V|ZO0>|ak;J0$M zEP4AnLrstM-{H7&=V{UAfqccfKUZ}e#@5_hfDfh7UzQ9um`r6C7rvd{3>`W0+GXz~ zL|}_;2@y;&T^=ODw;#$^&m>bm5<))=7C$jl8|{P*ZJt~hDliClu5AxfR3P+BrCs=m z(^AVp;AF!T0wXwtlLl~5mDWw|U={agsp-OHtN%rm!KCwz0bRvHN*+_o_7 zAQ6=$D+Rv+U3R3rI9OCo3GR1-eF7P z=Df&OdS%12HD37I^n=8Y%Z3J{t=&$?H}c-tRHUatI;!7o^Up?la>F;1M?JN29~~xP zN-0Pz#}&;E(;G~7#K*a4yRTMffXK=CFhFTFaVT}^o2+r+a8LwSMa52C>YeR91%y?n z^myX2{X(E!@>itEo2Q{y7Z#yBm+y@A8+G1|LP$AmWyv}>HG=Pc8FN6-qq1~R42Pq{ zOb!npnK@VdXCQe;0wxxcc?Ja?MLbg7WW2$MVG)hUs&*xDlFZAv>?N1NU#?HvCA+z9_z2+r?Tf}+vZHnF>zZN@uS2wbD2$ysBTkt_#Q%k|5B;Hw ztV+~(C&y?z8@8BN6?WF{hlaS-%xBV3mXrotz@f)PqPs<$ojm+0WkE&2R4frndY5{F zsw?g|Ye^0S>|L=7)7daklQU7H-?+>(s1dxRrG+yvEhm?@er~Y&kp0!YWB+4m1<( zlm6_s8gbMXv9K)Xe8%0yhNu9a28e{<*RQ0}XF>E)Uw9cX7#Lbp3h)!u)v_7bi-1L2 zf|cLCrBU@I)0mB$QZM>}z<2O>JnrQ7&-ph(8j_as zoP5>WjAo6U9V;`zm}SBy)hb-(SGqPgq4Wiy_W#qTPX%=2ZL-YF8Hb+YPoryMqc>I$ z1&h*bZ3`psJvL-=F%~u!^p{Y6Mt9ZWgL}|0vtJQOT0`sa5stlIGh;)ZOC{>M&lnWN@9jmWV9bD|*-!HJTD*S<#Z+yt* zOy<9*O9%UfIv8(f7oyc2o8uNf#}Pcq!Ny{iIi@JTAY1qIJxE5{N&0V864{SLFx7dL z64w?kv<&m($42(?Ni6;*x&JE`t^E!8I`gJO4uxT-?Q)kSAXg=WQiu8A#3tQAt}k~i zGeA}&%IK;HfR=CEA}akb{Uk=lVT2P%r$&nK&uV_ckCDubj87j3Q0m_86RzwtxLj%b ze95J;&Ojd?t=60K#wN{fR<^zKje2ZHq~Jk~oqiWEtj>0JEX?|_7*RRqd&9NzE*FNQ zdl@x=A{leX5xBwr3EYGA6KR|T_Yaxj<=;{Q$L-Oe<(V?!AcB!Qf-p0L7mb(_h^nf^ z0ClNR<(OSZ*Fi%$G`9Rwsi|yqUzDOmkJVDu#?_N3~-s+b({bZJ$1($lscnDXp!#bedFkHHk+;ZxUBN;7PC1` zGPBhw&qC_t5jnE(mZ_nEy1M#$zAj=X;(Mj!)p)5eG6erLR{&)>hDtEu(T4j$W3T(& zyQI&8->$D<50*(zVKy2N!b&U^OXbc;jR;-$9dYemv5m&i!>EkwG6oCGY1ZeqM>u{a zhvvoQAjl*6W$}24=76I`>@1~cR%X7Pt#5_hA_(m!_C7F^ab*b^3nvjq6pWaDeuy~N z4y`cGoe;mMiY1_U=2}5hAV4Ik$+M=Q*^$0}8Zqc_OeL$7>5Bny35sVA{%LasItQ1( zafdr$@GD169@P?3dy!DmU)H!%1u5tkSW$L@Oeu#gw&FH)kqf7eS+GvLgq*saRy11dLfxe>t$b2XdbH?)AlvEsf|i6mOF9fo`na zCAbS2Jl9w)t>(~l)Y`Zmd)mj}=l;U}{m-4uESSjcMFcM$mD%<}2XQ}<{^2`j(u_Nj zs+a<`1uK5ubp6c^@7}$8=GTHj4ipp)+pHrcQMU~B>Q|=Q}dugWdOuT(x7cR&i&-vO>nWH>RPY#dwrxs(DxBjsOxc%xQUlx29s*)i=EO^%|$hGQVvl&ki}WXI^EC|)}R1h^yhp4;ULo5dByKu z@4+DNi|?6pUqFm(Ri0Vb_;2j!Zu1vx(Mcb^V9il%3Omz)u9cYCD-Lx@D!P>3y>}0R z>4(xTljd$P#T__VsFM;*r$)3y9XqI@;@}0pE`zu|{>#VfR^rE(BjY(qzs#hyxcpV{ zSjt2(__6%nh5MkH1yjeIAUf(~Oof325k%E@}GjSX-Ezb`&y(g8HJf;)hTli}=P|KRiQ<`KKt_ z(AN`@j6!gNFef(Rj{eEMFLhWuC)#oq+!Yo18=3k{K zmMMM+j^@Q2q;LGw9<{Vh?F!#r*wmKgWi-dG}0x_3?(tt5F#QC(hW*8 zbPO@nyY-Ie{r#WkeBuk|oY`mXwb#1VTGxfFOMr(CjYZTb5Z~jh$AZ4` zsPI1JOifkO$>~Hke9XR zf+8eRfx#;enwtSO&Of*atZx3w+#Aa*`jaZ(J&D^62cj zy}Pr%ZqGep{c-oDN$bhr<&@q#|0fF1T=6Y9um>H8;Fg zar7`XCBWwS*|$0A2Ii1Awx>9y=^uR>_2^rQZ3IPS)4(a5?NxgRMt!Dw`Wh@G}nK3-vJ_X_eQ!0819gAobD%-_Kd-Ue&7d;=H zG*0b4wtQQY+cw{NY!bJ;$I=rs$57a-&YO0<{rMqdzS}ZxDMFGw+#0e}vY43w2xzYE z&C!D*{7!p92Ok=XDv$)JE5I8JHBwTO15BM>dvFjzV`S-{l6R<&^Anv)8TK}&%%<^j z;qTUCw>3QNNi24p6i!%zf2&eky>7D}^kVMtZU(T#CNF$)L5@#y5CRvMNpRd&^T1fu z5bN=Q7i_F%BfiNBql|WwUubO%%RvU7wang5`mr^;#(p335DTmCKc$Di4~XDYbLK4i z8`CvxBEHF&!9H7EUG?>*GAzVeW5T`?F4p%LQos1ISnpTupdic8NQZiGMf9)--gD0tK4R zjvVgta~T(P)Vq6u0PvWApl6%Spyp$A1s7Ch*8anTU9WJ48@$lcKG4^t0}&!8Aql@Z zGfOUQEAZOOU%z`rl@7~`Qkk^eYZs*F2&cfvNI{V}AAL88*Q3k9;AmOe+32<)SaDuI ztxiGU+W8&OyQ3>50xqpT(+OtB3uz{?Y1vnc5(vWWVP95V02oL#XZS+&Lelul-Yh{H zDc4t8kd)DMlkwywo`V((QzTOa^<-H+Ty0sk+}6x)Vbh98sDc2&GA|aDzazGkmOXc( zr-frb_6vge>-v0k?N?pTJG-&b7if7*aU;e(5@EIO+lOw)pK21@*-_$;zw6P*ZBB}i znWK(Z4E$$)W-O-ag_GQ`k?o`uoM2j=h)BZAExv4iU8zQC(*{7dYRXY$p_5EJQz!ik zv8?->E0}ZtYBj%i^QFb1*e5p;GNr8Mc7L)E-An4+ zVvDA-dR+w9#oMAVD)o=i-M6fl7kCnq(i3d=XUwZ$U#c#ey3(diqj3^p`Q_1EbEw5_ z5)zU-r!#p85~;0U;nN=s!*}1ry2vTDP_d_8$Zlaiu$huDQZ&D%zSMed9{a-JkbT#2 zl@K99&;IDtlnZ#>i}m*{mqJ7uYUhCC*(; z?V99``94w5rWz&w|Lo}IZ#()4l+l31;8l)-sB;kTI0jvcO51}|9qNqtGiI6sI2tMF#a0lhqoASXdP2zj0k*L>Ica;T3r62>DCoESh3m^#xIJ58Zmvh@fO9O>8 zuT|d@XefNH8vg7j==G@zye08QKZ6a_C^CG!TIcdAN}KF-SQDYStjr>yJNxdfr5OKA zp(cY;fDBuW#3bIwf-n=7Mp7zdJ&!LGAU@t%wl@3vsVdf-I>k^U=R_rWcV(}V7|#?$ z`vRUKChfx@P1FUOBt(8@HeBkZe88`y0cv2DCpR@||>o zN=HNJyMeYx82~nSP@vdf&4Sa0xBa|#>)iOlkKHmTVJ;Gc4~$peax}OCFYQdeEXRP< zRt?*C$0R5)LoiUT+QItx$Nrx-8git{_C>(rKp)_`-=_kFt+}p@*4L^&%!@xBWpJZG zNM*Zln>_7-FO|5N=h=+>!;xlWtf*E+UY0to8TVx3%7}1&Ra(3tB5Evgh^M1q+?UWM z7N73ympz@#fbKObi_7N6V)j>$cm%Tsa*uLD3j_;@c4YTG`Q`b(AOgyX97o!oLw8-*2ij@;gopAX0CNpj+f{C>q2{^_gzGp-f>b_lQ2x-&VHbv>Bcs-G z;h*kCsrMR@3&;*%FCN#DD0Q?|)qIeYfe59X=Z9HGGj zYZmOpR?@KU{FyZ`(UuLhfjbGuQg`plWBxx#JL~uCFh?U0%mcD~soWZ$7tv<{N84@1 zFuSX>6Qf;yOS|XZ?{2<$ar@!Z@Q)_IacO_r&{i3&IA$kOPbDPhfAEs^Nz;@R<*oKV z_~AHt@gl;445&RANUy5WY_5v;zX&%qSyQeP7zM0Uu}w1(Uop_e zkSK};dOq$~Ee$7WhVz6lmabVsCp!x*83>h@V3f^tS#D;TlDI%x1s#QtYZd-I&#RpP z`^)2&D^jvg$6FLyrBl_qyAzgOGuQ7k%15`3jE+{-#I2*aN$w|B?o-=mh_)xRnuSZf*qJ{QK%^j^_?&U?yj!BCO>#-K97vB)96jy6 z$Q|)m4(_G|oRvo(8r78;dJIRK`Yh!!$yF}n-&7gZXwpj4pLnc&e!aFnIh8y$)d{A% zCJGdV4Mdki)NlA*#$0#T`qEgvl9A8hv2}nxUrvewo#5YeyNy}q`ynRu$)fqkQRrL? zrET+%1I87CCKas(8dOW%End}bz3Uaceq&>Y{3inlu)MA-? zvig-{$B~rlyIlRdiRV{u-X8Uw$F&?-%y9FhT#xTEsnixl>Ul@GS90>v6Nl;3K~h?e z+^EW_9Wq1_j-gTAv4zK**OwhN4rDdV%iTjaGiI9oZnGx|_~lF^ErrFKcl|yzyDvC) zdxY|RAIB3FiJvNab(w9em9$@XKl`n?f1V4>#zH4yH2}ToYe0zwea*uZx#jPBReeyn z5<;1Q@HrrfromMxFqAQ1%;F#stLqt_-hdq(F1pOO8wQeivx;R2jJ2I!RAr=&7p<)y z%`-Kup}1K4A`HbZC2uZn%H__Wk z$F&M+-yvGjPZLNppK)FkA46@DrvlIAX>VNa#HWk5IQZ%r*Y8V(txjscA_b?u5vSja zfy1?YM{0XE|AyQrJ{=pOdQ^F3Q9@iktKU>Nl5j-*tQx!v7q#Mvu^w4E;vJ>gS(-8u z?VXCSr?sm#k+Y=r{oMBw*_YC@ogDHCtaL1cwD^}d#8J#&%c^KO&#^UT2OPtjYd;OT z`to~Tug6_gOe}fFwzO;0rgDDUYH#YhH@eGbw)3dsU5xLz3cdb{=nsi*)VYGQlsl2= zTBb%nstH( zsFaAa*aqq;9w~0xfW>|?YYkN(!w(m>zIGSOLoa9BKnx<5Wga6H+P`IIYURKj6;!Eh|Po9WpJKvm(GEbT$`Io8Zyq6QSl3Ec~hrl4s zPRSV5`Kcfk6~5&5CalcQ$BmnhtAVigyjW6rL;-%BJ>w!^9Xya~4Ea{?x^8DdzY?oA zF*Wqv^wVN$dtg6R1_ZO;V|%2zdM2)5)bz%Eo~obUfxSjgGrBIRF*OxX;O;>qlXt9` zD!7TMvD0@)&+<8Kr{n4awOOOSs#N>UzVEJ8bYS*7w(IDRG}X4dp4K-wYTOmkZWZZs>+L!cU z9c=){@qr9-5j74Dc<+JjVZ(=|6O276P(v!q&u8U#{yD6-k>rCHBYTJrrnuo(qyEKZ z$JXaft!3+?{YZHa_KkT92Ln!|K-PdUOJV8#+rsgCx#<=PD4x-()!lbB86^=w7f-hA zPEy;Ej;O*O`}>*w&>M{8Yc7!tON<%P^No6X4`N(A7KctR81QfS(P%ySK}zm)`*_C} z^J8QPqeXwY>$_Tz(AI1Ar;;Y-I8L*uT)x!Ca`m&Afi^2qVAmz0w9Lk$i)hy?${-Io zYTheBXJXI3HB(x`n_IgjrSmHCVqWOi2S;cs$yuZ|_PIesSvFud87adTO!wD%wyH^b6`mL7@~ z*y=EE*zuAPwPmtA*Wvw|I?Z3Sw1h1JO=9=uxzQ)**4LOuQa*@O=2dFD

    6J zikKiNdDg7q)ye={mf?ve@y+i>0vd)z=`BQ}bB5_nxeY&;Di=zP#?T02A=k>%@ zr|^UG9r_Xyfg$5>OLJ1nHNyeVDNl*9v+VF@#OAu@6rL>cq$r4rhLn|Y@86B@BD!-@ zz@6KAgQ7pUH{0k(VqOCViB0Y3bMd-oT1JdIbLa7RtXShn34Oe3rV$Kq}hbGR#jXiT^Xn`2dMhxJ;w$26LkTl$>LdMhcYNs{{bccj^R+Zhq8y&C2|C> zu?|UdFhy1x=60&LOx{68zAs4O{K^@p{Y4==5wr<2tEiownXDEoPGokPr?0Di?ixCF ztB8H~3K$VEKAC|uTBra z4A(>$&P%h|$H!uU-%Ys>KU{%Hh6%f}*w{BGuaUjBw(9{jHU;6em(Zq}>DqCg*sv3k zZNRT4#>6yCJhwv<0usl+t^m?sf0up%%F?(L&puKWD3J>tPKUr$RSk2c@1AVZTnJB3 z$=a3TR5_HiaiY)0-?03wmGV!Tp<_EEx555HY~>_Q*O@7?Y0ibt{S=rJtHP<;T`k%b zQgY`$#hVv1B?9GSPJ1Jb={{+&B)Utkge44?1Nk`E!3n`#=m~Q-qjDq$Zq<^b()Uo_ z*`p&vZ%wqS9Hvkot4_N@PQK5ofWCd~u|N9^^J?N;AiBn;pR8D8C;!z|G1Oi~P>gSQp1(<11LM zj4O`-A>4s?xru{kMK={H65=YnDX>17!BPG?LRVWVbKZmlkjc2#!==11F*JlvpB@Gu z3m5Q>?s@{^qHA(5CDi$TGk@slq#P~&%lyHM)uw*!t!8$0^)CW+efFfjpMWIp76Fr{ zZ;y1*I5B%KHBU&#;^C%)JuTm-T0UQ|fCG=LnU7h5UQR0QdM)7RSp9 zDYoo=aOW2P7^4VPZE;iH6CdG~C#78I%W-OP?~+8#}XVr=!hoN74c!qoxMGOW*8f z6&qFyPMtCX#pcWX!aEl{2`zIjE8h=+nz&9Xc}IrP^Q-GY{`U@4L5sKRI>MFC7e+=S`Bs zFEf}A&T}MJfr1S*Hv4tfjy4AG@nL)+Spn|H5SR$A<#L`JVtxHD(@@9|5Ae6QFa2+G z`1Q0(ZM87^H(ha#WCi}=P~kcoLZc^URf4`^0Yej1qh6PXcdx$?|CsvcTk^dgAau|4 z$T_KZ!zd+M^5I1X+;gMW%wd)@;mq+qgqTg2KB)&m{h!uau3ynW^*xj^b6^)Q9|$3= z;}^c+UrttPJR3PG*Z#fa6_Xs zY?5slW7B6F&}ktsb2+(()3xO58&%qNDHWGxin+_}62F(Spi&E|e0$&lDX8Ae1y5zja)VA!{=;!2teDm1Xo%;D0>H z^?&z)wVm(a_%Zmk8Zg+Ngai@(U1&SM0G|jEuT+!=?B$gWND>CN?IMH(cV??74|iw} zd9*T!A2_)xA>9>hlD(bVLa0IF@gat@K$+SmnT&#>=R3T>~ciOUcdOcJAdzGiesIBziEQ|cdAYbuTf&WYVW4^G1* zBn~x>k%d$RsEv)#pdc%8*?-LBIl*7isf~@z!r~+wmq#=8~t2Dv{7gUkylH2EAB~dGZMEHe4S(4-nAMGc~Ed=WqDx^ z1=uU@rhKt@7B@`H1!JmmVG-Mo`oIvmPJ+vM8eU^>q=d=q@w&DukcYHpK<( zXr(7CBK_yjY|s%e+S`1x?vC-iMiZohefD8d*dzugw2uq_-qn{s&xzIWAX2o0V5;ta z+!oA0eUPQ6VD|u%`;P4TAHh5OE@@rv9*ztS?@?@>KOIjJBmcD+U%KZvQ^fL_0$bxy z8!h#ovCT{V$0Kd(724@3DX@z>|HnEIbzTQx6k&PqG#7KsQRka-sqS+?C@oW-#TA%c zxWU&nzp8~?ueuHdCD}@t=dUK3ty`+st!VKHJHCETn%T+p2HxB*JuR*Hf85$_BN22` zsyt2>aa{{_4A^d7>rPOtvZ&JWW+B2Y^Q;Q}`;~3(-5AaxzG}HTh6cx~>ZsxSbzN+$ zwy9$U8k?JIXlob!_q{<^$ET_)o(K0a4QEAs_(OARo!WQdA^{k? z-|Jl5dyF?Dcp1>LcKE3(e!fn^$hpM_Weh08C9SPtM|A&ZV~#C=5DXU=Hw1%`7ZL4& zsI4g@1m=;+$@VU*+$jyvRb1ANe)XUS81jUrZ8GsGST2?6Wt~$v^)tZIk^-b+OG-*c z{Dc0ni?@uglkZ*ot}6lQgeOqYZ1KN{kjg<17itw03}x9k zHY@M#oh7jikFGznwkwZILar_T{(`WA$3LbL@bR>u$H+MCQ(AT8Z|usWVs zIQ;8JgfA_2mY}PCN`V&qv+-6f9)ty=mgBy;Hj$khpKFdv$UoczV!_$v zj2!oR2lkz+1;$N3CG~C@B@iC9g}@GEA{U(NB_t%A>-Enb{-5K^HZ~x#38cpouoyfrSdqt)W4?~%haZxL*53z3uC?*nP4Tj#l zLlcx*e%_b)HX~E0!7e0I>8KPZoSIL8O5yuy-a7we zY-9c^e4x?Gj(vP~wh1sx{@3ahtLW*arKk5UJEwq3K%iuQy$dj)7F{QW-l$vHu5M_3 zO4fPtXV0?;em$gM>dV{P>$Weke6a%SUuzn9sm3g<!xL9KEz)@{doPcyv2E?Za^7#Q05l+IAmZ`YoBw^X)bNS)zd;|gj~^d8u8?JznXV8}b5Yz>@3|Mfx9H}E z#x%J4y>Y6|vEuN14K~GlL&bLiz9ZBov0~7!Oz|osjpi2TWQ--RU;YIXTtRNK=?zU< zPy2sRfVAIEiMc2g7S#Ow$Iq3cl_3U9)t1qg%Zo2C9P%KM%nI#B+w~F@amU(?8^HM= zzB1fOU74(syp5IA*6eQ{mh((Dw#CRvAgX!jl0c&%&+hoh!lgbcDr)JW#Q#_^LqjSm zs{Yl3$cP9tZk=N==Zj03BR+jOYBMzCKq)v>EKdxWZCA{$R)e26x-kFREx-~Y@)Ew` zbH#U2adrM|JILrRZxJl1Lu(n*|KL#1OmBSs^?tuE4qF5X38=3R7 zyd$&ClD7-;gLxTe^iCaMOM0hgk7^0lu{XtI!F!)=>d?i?&VL5%y>$Q65AK1?Sj-Ot z>C{R>$77Q1e1**1le_MEI)ZqxGrc4RLiK!mxC-)lzM5YisKh>HC+u z2UvUE->Ki&!uoKb>9>dq1Og4NgM)Z7(#VbZD|9>buy9sS4U8vZvL%}+DRo$h()!#h zD}*ZjS_=!`?Z4R4`QybMME{c1yI|#-ota&7N#*Fig-G(rrlsYwhMg5kb+fNSNn-&; z19Ec1G7;a+xo2l*=UW2GFLwX=^4ApGeg)eI&$FYga6*>ezP=<@9p|54ZdqDd3bPdN zZcW=-DKh5=@CKmJ6?CubI%h(jLbKm1XRfwl0461c^ss&0b`tR;q*Zc zOwCWQiiNhNo#;hddhkCD*O+I!uO@=6p|4z6UVQZugsi^=UKkl)WWNDQ3scd5%hq1I z_MP>23I~m*76;8|X(=gQo}QpWjgQx_{!guRBLp{==gX1`J5^1j|dOb_r??w4JYZgcR z$jpQGJ~K=>22Ud)m8{zAZhT*2w;2fg{>S^k;qaiKAh0Lu8ybLeJre^1Fjc)SGZ3m2 zzf#N>znc5P92XBc-k}~WPI?D2%_=krYsI@aOEK&;XN8t?+D0Fmnm}=6>d$M9=NvEA zuH4PCi#@BDuJ&q#5i}H8FA9AR zyvLu8UH}*1przANPP4EzZe8AcMiCXoUL`3bYKuKNL+YnQk$m*ne01Zso|wqI&c3r2 zfD}P}LcIxH+<#EsbzIJN=cMLdi{<`Kmb1~6N3>Xt#2r{}7B>#9W&V-Z)MHHd_s3Jr zFNxUI+?vI8I=YQ^PS2f>J1FMQ6|&Brvb04HWb)M~sWaP5Yle0Z-Y0$SA?WHal^HNA z#WqpgDVkkJUMSn5WF&i)%?wfd;^*scKdotIMdGw@)9dl~|N6z@fP-A@$rD-H*a&ig z?VoL-Fc_@q?2wBn%(t--IJ)7iU`H=tX>mb8A1=-Uj(0zC^MRgXyJdEbf-)V~#oRTEJ{{f8%B~#uVV=X# zl!~=XNw$<;T7Y_JK=-VaFiEpv58laniG5#|E2vF&)ff7~;XeKzkfSY}L-adT7u%M` z+Fqnv?V?tULpytm_t@FUJ8OSEnR4xCZ$?yPV#V~7qgySU4ugFynR;#@83B%_S^Ko* zb@{>;jnvU3^(m_P-e|!3{#q}omawoWmg$@M<{hJ!G5>!+r)zJ^e`9Ebu(C2GB_(%t zL7$^dq{q%5H*VmxoA}806j2J>3=gjlv0ROT%NS`!#`@imga&o6mOk>D7T~YVcC2WSFU$|8sF>5FRk)d zKNC5~8r)`|T-h&P_+Bob*Lq>yR=6w9&vO0@qc$-(f5SuJ3|MJDF53nGE~p7u+p3S2v-xb<{f7hn|dzhnCxib28^lqD@T|HIHQlCD|ql$#O$j zxJ@37Uu~0}DaPJ9HPxSKSRbBtTl>7aM|-@jF*0H${tt+F?eqBW$fcpt>%d~UJK@{6 zZ=l(k<>^ze%k%o$+S+%ryd_%YNvi`_7hcQ>V9QtL<&}K@4$=q&k41#5i_6wfzI>q) z|H{e=cS=RE+@bLt|#6D71S&K{bZ5{C*+Z8JgCab$6)MCM&A&f}v6IAtRO=nX^@qqfCC z`FgB;EBVSc_n2p-Akra?pK`i-eK-EF<;qX;(5O)6sx&eGIDywux>THBGmYOHlv&4B zLr<4wzniY}$3{5>QV*NPOmuhd4XkiZ>X&wswvL#_Y2PP8is8n$o4TA-cZLsYn`sQ{ z8+*QZrvlh>AJfr6OigW4hPtPEOXQi-GI0&058aZ*YGgjI+zpGB$EG?R1IJk-j%K3} zyi*W16OLImZ|{4Ap*_99tf(LYz)KymPTkeTErS%0(gQ^N?~Cm>P`9CGXz0H_1Z1Lx zWn?;vii(DYhJH!=?aa?le1EeSbeDreKPI5i3tqV~z6|%+JR>1aA|Om4B}#PKUqb)* zk^J`U#M`%derFEU;{HtZL9QDUT>CQn?IZzLP^5dX!ph<;-(rXh-!!x}tA56ce-k{c~fQI;;Ilxq$OBIA*F`C>UkcdoBR!$g)AKHjA8-uwlQ zF5bBC5{3J`_)p`*O3gy@lveQdTujX^;0pOQ1>MmZeP+0c1yPz7QHd*2L@0I3{Cx9S z!Q$A!QhU7Hp-(|BIDrN%kytvJIpFdD zVMlg$wgJ#2HQ8Hu3fbo{x->k4q-|o*4P+NNngR~IR2NKXu%EYzP5`wA}kn+S(RYRzy!(c{et#)mA160zxQ#9(XUcPD^>E4c}t<9pFcv7@5y{zn=wI$(5HHxr{W<^n!f%OwBgdmKIJ^N6V|~w- zH0K5&Mb8|JLd4LpvvK_1 zkfQIyQc*U&hwxZHHIVD2JAjQ7^Txz?u{DY*B`W|f^j5&+CtV?i8yW>OCv$;<(sjEyjnXyV^V&yyQ z>r+hGQ|3~fJa~p2% zM+pg2(fRrl>aac8C&4nK8_{e*>H5RugLcdYfa{kO{7L)0g+*SrQ;;GAC+=a>!%pm9FLu{ zdN=vltnv$ze?I>;w%D*Jt#Dr~53go2zQQ*6d z=H(|aGOkEVQ-OTCZ|-v3eQUz5CZ%F;%6j--ZZ2W1!_&&jN9N{J8)E;gpD*3-^b3JP zp}^Pm^72YaPM(;UC^2dgV__-I%M0>#S|7^K*xC&GLt=p40GHM1@)_%VgY$C5 z`d3+LWTxw?rAd)ysJw}air1^~g%!p2D$DSF=Pg3&e|}0c<3KvA);3h20$>N3Q4(c_1O zh5~bV>F2SXyC!@J%U zI#I-#s`>fy^SxkKls!pLF7i4LNV!><8>dno+7da~ex4%9M zGBPl}H;~G=vCs}3$Pypu?-%pj{Q(FhBk1volE84sIbe5{(Ja3_7=e9Z(E{fx4zWbr-s)SN3dyB|i``0K#Qwm!gX`29bwJg=x|=i&mN zz^K7Dh52fKSxZ%Qpsg)I>ZJ(uy(a~cw8InlZ6F3O=zo4RnA;siNjLk<_r&`Rrv~cX zjX|V`(HsJb$|>K0MUqFC>&duz_N8WE6FX$^$W*aaS-_?b&vtqkl@6(REF}C60?zEe zeKO9xfmK94-g32@DS8`OS46(FFw~+~+iZZeSy@?IvobYxG|$(u#KtBF`zlL6q+s^? zyQ=ceIKMsFdXmu;=ryCe*vCHABv;Gzr)e2GO2Vb|x1MkA>hutHw<|23HdHQOUgsuZ zWjwHqYJFYpxR+p4Ts$4(-(!W|D4!B2zgNGG*m*R1F_h@eRAWCL)v;(jyUtCQ+n<<- zy3y9!4G(SdYY7Sc{t&TUxELQe*uvq*wJ|3oRJ07*)uYau|Cj26{6#be8r^)wMMMU@+DTueI} z2)m)c!3@rkNxPbNDiiaQr{B911Y=`Cl{X=Nq=dCS6mM!n_@9k@j`Qp2c^jwCesz9& z0K|tA+4L+dEtz6vLAgpgih}gPgXhi3>FJvr8^TuKAG=^X&B;djB(bmU2@H&y@R}$X6Al9I_12rtBEiN{hmfyWR;g?x-H4IFZ;Zr{`B>-M zH-D8aYnBFY#9B8Y9H1{RM)jAr^5x4p^dC1hyrDThKV6ND&9hwX7sbUJ5b{Yx8C$Jo zS@xg9rk;C=K6|$H`MO&7yBl2BuXii@G&idXdxV^w1li9i+?#DY?JZC&doK+5r0do{d<0rzY9oi;VI8{G6>vi48OI zUz)Hb3~iX1yoieVMkf#(8JU6L6FpPP@J7}tkO{di0Ch2`W?`45PV;!i(qxhP;L>z% z^MewdcTl2R?1ty@HCBUVrw2*3C$HG{5A)`^btBhir~8NE{OP83w5{gp;6iE1ROk{Y#zB0} z8ofHk%kuj%Dx*RVYRr31DTRyBH0+PxmXNI9##8IL^Av*kBog-Cvq1Natl5Q^-=4|b za?i-nOE!qvXij3eGMe0e^lC7~c@`3B4U=O-@Q;-yT3egh=hbAb$jHP8ECBBkJ6*ez zEZ$PvBq}(rZRajSNy!M!=XekURCP9e&H7RWlGsP}TgvME^!9eMX*rZtbzn_(H~zNx z&+~!5T90^Dbalz7saxLsd_hc19QM-W_3PL6_F(_C*w57M{b&so2I`62=Ha=ybXf!A zgQFJ+dp#_uxveo*nNbVC3@8k%Ru;+(tZ(CqiMgyuz5IHxyW2Zlm{5M=y*aT3LJB;p z2)ujY*!LPxJJbkUtqHONQ zUHD4c7OG6m9b|+0(VEwF<8~Bk)b>*qk9ACyi=)q389%=gKY!HN7|@7YVPsTSlavIO z$<538U>CG&Cw=|-#+4+(NZ9oRVcKB?s8;1ib_UKgnpBnZ4FLZ9WQl+>!(D+9CK?Yb zF~`PjBYmV?Ca}WQ{ksIH)MrlZd2%eB9T-jJlm+CQ zpREZjFx*moFSD(w&aLUuG$T219x!ds%Fjxmwy8k5d8tdB&-=MprVtu95=;124eCSM z6}y32R0FxDFWScX#UFl5M%cqMmNV+zvp(bH0eS)s=bMQ0SA2T~HM)b2(-v@9C3Xi- zaUyA%)pYT4w?c6oysLvl>pmc>NB{ERq~7h&Dk3WjLjs+tQS^Hd_eUb@sq+V=)R$3u zsB4D&#$R$OOfA(ykH2R_uyAss&Q+Uyq5&pUF3y2*-%-lorq{ZC8>Q}Ci9i&I_n4`{ zR33lqe~Iquw5)!7SibtR8+1uO8$#pgP z9J_DN2>PD9k&%&sKp=E`wowB@pSG5Lw51@*J({Q>2*u0d zw9}?$D$G1Oq956~&OOCC{}u6y^Cl}Lvir%vAn@o3LeWK`(e!d{(7rizgfJrO{7Zj+Nc z(};%RKo8M=kBE-K(#MZ?Sh#rzm9v;Fi_Fc#Q^|t$%bKG8b*dhJ{&uRch6W*Vad8HQ zW#B=1Ec?MGs;a2q5)csJ;^I={8hL*CBvv+04b&1~!S_z(czAdK%6Gt>(yJYekwWi; zx0|SnUU+wF)oC<;D}6h#3ZcGB6)m%czsMWnn{l3+uP{~$DQs3!%Jn`zqTqZzw>Bux zq@c|QggM)KQ;V2nRz~yK=w4B|*+h_Tesy+oF-#ASZ6zW4wt(?zyHsaf&VZIcte)q2V zh)*Sa$Tv_LMHt5yntP~f#bc_v1uGnT4;3SnEV{*4HJu>Va>dOs>Z$06GH=?L3)fEA zztzMy{0#C+eHa4Dwx%#Js5U|fV{Y_DYp6e3H&-7ya9Pca8RlnY-RMm@V^UI?#lGA1 zm{Qz-arY;SzkkxjIs4Y6i^awm-NrbzVT=DxcY^TcNno}8&o5|~m68-em#xF~9-zi1 zX!7*~jpVn(9r5q&`9`%4KoFSU!h#Y1u1>&+JWjSZmnnmUsLiMvhsj95sKl56dS8|X zN{_&gV~AyqF!SOqw=_2&1!jVq?{VoqoyM8*|pH;ANPZQzSci zb`0{sg8>ylw}X|wjj>X?76c%JzJ*QX z>EknyA?z^)^pgt|QaFvj)HyAFi=q@Z`TAZW5DSNd^CO_f7WF<{8!JuA$S_##|D{?6 zo4mmXV^A;OhS946=_}p`VX~m|5~WhqvAFJ3-$0>s+hDdq$bIvB86~UYrjwZGrpy@`SX3t4yH*@<)kC6uGZfhSP=>u}ha59t0V(T#ylWI3MEK7-E|nUqvn?~>4$CbNuu-TH2{IJ(p*z-IB=a2kKcbn``J zi}=&2d2zhOQWIjP|olsA$yfTN_EM~9R7ijDwFWx9z8er>QH|A+M51Q zJ2WpTiG%oS2l3bT*n5kM*uSi4ak0e}Z3Fscd`2_J8x|Hisps3j`76jikyFUS>B|cfQ7){+#Z6!h3FvaVKDSwF*U!HK5OB7~@Os-xtg3T6GD2(cCPi!!Bz= zm)mf~`E??h^=Y?PVc19L|#V1VTOcIPL1W*0;#wIS)uYWPt1xBYW&CHg-wY=OEaFxkg z19~q&cdFS>4F|#YGRWTt?b`stUuelL@#xst*hsO);ANPFYKx(UhDL#!t}Zx_LXUUm zK@_@L3_Js>9ZEa^IB+>_nboh+AxVgD)&t)u) zAc~s;ZzPmvWaI#)rv? zXF7Vqx_a@-`jiyj&xF+Kr>2iT)$VL;ikN)9hUeiiB<&vBXLj64I9QGEWFo&hF}=c|3;q0*w-BB_xQCw~&t> zjmW-xmvnwnzCD}fyCkdDhL1-LsEfrJpS)+*MDwh7L<_8~4IY|yjO0@ayNl8<8LO)Y z0*9LKn;Mae1DrNs_AC-dP#H!rHrN#=1eIHgC{Wr0t*P8I&~UHEwA`V=zM~_3^#9oU>ZmBY?rj?t0VNd#q#L9g6r{Twq@@J}qy+^8 zgpqDWq+2?MQW}QtAr%;EC~1cH4*2@Kzx6%u{fDwxYYz9E+IwHuwfEkfRf@$g{q_Px zQ4)d7+`t5zP1JxrLrPHkIKqX2i7B)|g`c0lrl#g5quEXL7gbeNvDk``kdVeop})e$ zroQU{H@)(TWnnMp>T9r(z#h(&1xTg}$Q?4A0jOUp(D`I{fzxC^RUMI>oE#mEkB8^H zF;xSEx!|gg0xAkVM@nky1>hiD-E##UHgabGoU8>7RNc&1iBO9%^y$HDu6I>SJiBQi z4*AB$hRCx_C8!9X(ZlyZ?S9ni1OlOJ^Q~010ZLWDq@^I4K9iRpOnlwfJq&a*stM5I zuKv1p^BOj^xY^8Sq9b2XDmz`#`3(>{&WfuBc-5N_SBVrr4 zmCyR0IJto;B6*cBtu4%KneE=*Rw)+Gmkbx^dH>l1!<^fxMKy?(|5Hjz3i@DE^jUxx z2Y^tV%GH8|pZ8wdd&F+pn+$OD0__t2WQK+LJdJ!F`zh%ZDInF#R4i(0YQozCAS#sx ztA6$8$Vf3orQh4Ye&UdeP%qiKy1E+HgCCVGdDG7NL~id1e5rr-*Rj8HnUr+pokKcV z+pCS9!cmY%p^^4;XQ_{753J9wqBzh*R$#?-ba#7ZQqa*Y1+9Rr$H&Km<`4kUO&+-q(2?H@o%L8ZZ`@F3`&M8D()z)0(CqKs-8R5D4@(+yD=GDD z&$mOYVuOO(zZE|cc(dH(cO7lDJgMz{yzi2>au&vap5swdvuDTd6oY|9CKNq|Zxs53 zH+R3V;H%5#ec&P_gEdE<^0{s3J85}%dfH6Se}F(Bqb1LkC)v3nUqRl9)0C~;F3?D2=gO$iH_W?thucrUTzBex5uv(mjgd{v?>(H2OZ(;T%xo23 z$gzMjETgto*LgWJv##z0+>Zf9YRbx=lakOVn^K@2=it~ll`^xiur~6`5(ZLFKxj%x zNFaEzu#O%%MOn}G6A*xoBZz~4gjH8gI4?2Tzkde&|LVrYo&SHIunY3+|2|(HTwrhi z_u;yBalGC&iWQx`C_1UzNMDW8tw~l_lhe>sQ`0k0OV8Vv8{yU`B#h|$nc z(A3cx;8@pwKU*`_K4uqIWMwKXe!(GtO1!WVP;> zz^cosiL?HEWG<$oqNAasqhY8bqan#yd{U5CUhp-e;6!a@Nh#BOVZJ@Gk6a|C?)-4& zG&|CbT$G^kTzP9ldLqFNsfF`(| z!E4>wqTW!=s~`$Jt#~(mRx~A7dtkw18`Ord>2HE}tg616;;|V9+m+|Y4FI~>nB_I- zP(>vb{`k1&CI$wS&ruBhl^1Z-WkldG3>ScH*$zsu^8lN;2oiXzl0HX7Y%g?1wTHO_ zW0=$i>@fj%R(18EfvE_n3SPnvY-c7HV#ktT{Z5&6n>_T<7+01L>j?oOTc`UKwB(;T zI%l|8lBPyU=U=M%Q{Qvd4Gau25wCS}5zaE}*(U|SUhAr={sv?e(9ksilba zAvmzzsiK8TSxo&w(mLf0fH@u!e{alN(B}#`L_D zeB4Ll}U3DZPo$hG3wI*Mjmyxc$ zC|uEi6z1yk`Za2*v$D+2GP2N)u1CP7w4}1K{At-!2WC7@=Y-sZ;9P!T@Q-o#$5zl2 zg8s05`{qewOh?l@{c(q`7qpz*+_}Y$KsC#wK;`x>Jv|*yK}&V4sJFMbKB*sk?dM$J z-X{a+{$tzQw{JfI8#@R}esasF3Qd4p;z18!CwOhgzZIqdhn|e@^^@u}aP)!#0|R;A z0G;AtATFQ?hlVz-fvKqWJwF5Y5eJC}xOIjCQ$^{r)vp87RNQ%#<(K?!L7q9wijPlI zFE-{K*b71DVd=g>R>stl&JpW|03}2r4v%Bui+jIelQ%Nj-Mn=XQ z3i^t;jVSAD|MkC!fOs1p|H7B^M(>a0Yq1$MzEX)4A8Pm zy#XIS0Mh&I`Pr%G!3Ng7ds85M0&Ca(-=Ymf`j@`W3SfwATV!PhL;r$?L?tHZg5SS;G# z-U(iZz9!{nqdFHnJO1fivhxqH)ufW_# zPASXE-qjH&d({DlQ+IbWVSz2wYx6PjDiM#my83Pd#qox5X#w{*N5NO_>#S;a#P)pG zMrYXdN45g8I>|eWKev`;WW3mcrdKjV6>H^Xbqwcdds^N4&sl!AiFJ9~3Dk(Gs4*26 z6tU#zFat*Yv>gKc6*j}Vj;^j+u*(3j36>UgF#_3Y#rie4;3z+M;0P8C?1A-(GRup2 zbgUSJwc95kBMLYoj=x#}WsE~AkYUu=P*haJIt`ngbXys0?JaU%QtV>!A0Y}D+q)K1LXCKNr;1|ZmV`13#00U2A-4E_r5v0oLi@OR14N%yJx|@4~CGy{F4=Qt1(bC7rNCx!yJ&MMQg`q^b?Z z@$)^J&|XB4$@En#)|dnm8Is4fMv-0jUIb#{?A6}zzeN`ydEBKYxpPldLXioPkuDa5&UwjQ#(55DNhQc6Ou3;TY*=HL=a$btZfKx|La z$b$&s8TeHj}YKOm+Aok(7F9Lw~)@1z5%oIij51Y#t^8mHE4nYF(5Kq3gdtZDG8 zqn)Kr$bWtS-e`VKPLujyGd&0W!xe}XO?l?!o@z*OO8kEy*D_6ob7oH(IWHn^T<_e z0K(ea+rKiSj_G<8p&j)(*x+}L;_xzK5pK!HdVFTVhXkPkdjpgYI%%mT@ENG-DFlCK z+>Bn96AOMuZ%+9^RhIq%CJnBOQ%!`1lqC~dQIk|~g0Gsd7R~3j*@K*bMa3Cu%^FW1 ze;Qqwa3~o256vj(x43rHDkJD{hL&@@w`0dy+vWZ&Hk}5rg$4@=L=)h9jXQq zZN6z1m%RdNH>^uyyPc$pdV6@)Kx=HXvBKxPC-54M*)8(M`t`A($zE=T-7%@3L{P7a z4!ORs#SP&N`;t}hHV(h+8|2`ur;0NC&>^5$IoCv!v#vL$Y^WV8#kyu|swgMKpUNDN z4zT8PowdQ1^AQe3t}DeE4isI*Z z#0`2RP6s64y{N{#)y8~n1gkyuCfzI=*GN{!JAaHmr)BSdOvuG9b*J{MdVu~;jsXfR=}E4>m}%^8UIhz95nK z>Xk2dVq`Tm3{KN_cj~j*dJX%6rG@82eZ#aGBV2onC@FfzhSG0(g=E~(eG!sTu1y=$ zA>YGnOR6ux=RWSTVLw)=BvpJ>9^3G`0Qw$kf%BU)#E;BJybj~u4jJFqgASSNSR--;5pZciNS-oRH^Jn;~-UA2}fQzTn~A2b&h5CU-P~zECuaP&0d` zXh0@YjYW)wZ+i1%;2V%Nm6rtI&*x!{dLtT)cNyay76Kit6}EZhRvQ-^@ng;yfBh;8 zqyukw9U? zkt*BCxkvfq={nx0Cq#$$$ePL4h*w1mZ2G>Z)M0R7h%2Oc-QxGa8GcN?wl-U0!-hP} z?d+Xs?kx;TAG=aT_Wrc3_c0IKE7Wi0y=FU7A=6CQ=2IA}umZE4OvHzJr|RRxUnV8v z0Oy&V^##`nh|o?>jyNU%yL1$nOBXJi$;ig$yZw7)M6V_(J<8*FpS$WCL2$;khIh0K z!P?Z^Je&O!^#{u<&BzlnZ~DG|z+;NBe7|on>$>4{*h&Xxw$-9840a^4(S|-u z+QhT>e%XEUIew3U0T#``pDiDu6%(nI*-roOZ-~={89XR4 zY0umWRU&YYLEYaCm*xCDoN5PB0*}$AIE?|^>|?X6pU|00{lSuL$VB%%i?qDBxG?~X zg@lIY=jT`89Q-{sTXdJX9A9(Q^GZug0VoCpUp5_|DPm(|@g)YtN^)cBm*#?3{NX9kQ)nAbx z{Pj2I`Yu(Hiy*57_?ZA^9tUy*EUbI2EgE;U8(7SWH~p(&A#FZbjGNU}P43B4MJ= zBu;86GwQ_D+_+)`XYGO5?b9ez>}frv*){GdDB+g$o`!FLSXzGm8gN{?ySmEz{=Vv! zEz^s*^_8PKBgo&i^6)tRX54JX@km0IGHpCJv@c|xD5_M%m>cBXo%Usg)a`op51W^_ zJlmnj<5Cdh=LZf__C?aA+21+7B6_I-j+Mw%^*-GCx%IPoW@ZLoB7&NRvvOI~MPmGj z?qNFR2W3v_2qS(LJ_UOd8CBG?m@F$#Pc-0*KubCs=sy)F<{+11K>r-|JP>V>0g37K zbV`<01YHus@`Z%weWb+#>5cVbNbh?}4jkT(o)F5}%XmDko zJ!=A(R=Lu>zr$=9Tq-qzd@}|(g_~o=Zb03YKBx@~ZAbB`1#mAbysi+_t7 z1@h3XXBaSTcuphPcsZ!<_#sUifv=R8lhZFE^Y=BHcrWMoL_<>(fW{SWo0ltP_0pEd z<7ZX|s;^emRSIMu!eu-2A2%|}J{w4E%>~h4mBZ|)h|>Q*2=JQ&0t34ioqCi81>N_s zgb(g_tAZ3l^NbKh-|3pACt^*Oz1K`i`1tr9LX+?Q=i^)NFZUl3bgBkqAW3W{abi5g zpjIKmD0gSb737!paGD$U`MHjRatIE($F8&EdT_L^L zQCyZ#5pP)q1%lugC{^m*DVtZ%;1!f9&~H>=AUI2S@GGu@5AEY!n|ti^yGmVu{=h*_IH#?Em||C{P?^AeCAUybE0WrrLq}R#`XYHe z;V%b#CGqlvO@M0;sC6Tgl9G;hRAuOj8yMUd5u$ohz7kzO5?kPq$zI`4wZCN}j51+Y z!L@=X+r_xfX;1cj{$uMe;Ecz`#o0jl|9-hI43}T75X4;oTIp+R3#7i2;5V?v?`NeQ zdt8+fQnT-R(sesHI3~TXiQAL&n+O58t>tx;h1rguWGj$FNOKwg_yBGX>wk|4^piO` z^gx9WLd_z;*9#FDj10;R@n6kAl}1nfZuA<=WzNACh_{*;2%x?}8C=)*;U>e8_SH}K zk1ezA@;ky5CS5^>Wz)N>pnSuKi1Xi9ex`Q0pn4#h#lCwNq|t+$vIc#FoUQ*T*F%nh zzshT#&Hbw1owBNR@b&`yOJ!fbtQH@CR73X0WhKWZ>8MaM;D&hKONmZK>$q_OLx+(O z0w#<(vLAl}fD%ZqG;ieJKlU!`a`#oHiuyLZ-8M!lA-?ZUy+>Pa{dYH$N=q_bXC`ywu*g$yg($4MdG8@a)w^U}17lIr zEi>Y6iu!o<9%CbgmgFxC+&hI#oa|en-u-Q@hQ5McKKCK*G%~8poG4zaZ5%J2La?BF zn{_6@SQZ0Y**TC-M>BU}Z~r`7RR4KT5pMy98KK0lW=`+V&M?)SnBMq;_l!9Q?VobB=uLllQ5u z>}@ueXS(}*2kPRRCwoPt;m1u=#;eejJi@%net(?xwS0(Rrlv{G$ai@IYvAphMpKYCHQEYis&LgK_0f3o8A>9Di%y4`A3K`=1IZKI^kOLp4#(cK~Cn9j`yp4h}n zGM}hvW1TojA|h_maPpM9J`}L@7DGz)6b?5Jn{w1#mB7W9P>ueaAwew>1zgScpN=)r zpLr@>i23X$Ldru2R*LbDVrp?t3vtUt`+~=nl)0SS3HMX9vJCp`-DM5UTUsr5zGCl_ zarewj6gpR>l6eOuO$xYdu*bXeO`~7jPHSY7@R2y6IP6oJ)hfe$gP)p+(fd4HWD1YahDO94*^yYn?`^`U7hYkn}Av$k8iaTH>kUU3Sozay3-kPmnEUgeLLBV*Z|n~ZOJ2xs z9TARVKMyLaFf>E9>LhxN4wf)Mx=k(X#UkMS+7@Q=Ic$8)_;R&aiT9JXH4jorRPoi@ zr{9u0SI{%adCH>;!R9L&3o$7o$XxJtW)8b=UBVu8yYXGCKDK55m*>E~BP_B1v#YMW zx?nT@%-hwgSGg42_E(79k1{tP!{?DE#f_vD^zr#*S-8NImnD{~q&*T3Oz}ktbG^&Y zmWs%DVX-V|{9LtHO9Y2|qYdWR-Ml+lS@O+UH~Foq+xx&@4&KQP=Gip{mNrtnkOrzL z53I;8ENQ=IjNHqdJ~PQ5uXjf(6-~r6EDleDiz+duy5!Mjd-10)GGg@jHE%ZYT+bqU zr}My;)2>HbH+IT>DnxzYlDK{7L^WR`$HS{nf0)mJsiD_5<+jNi**<_5Y7WQiuMKyU zwZG|HMG7jd>>P4QBPP?lW9BV@2WSzhGhy2JV_Uks*?w|} z5>qXCG3vZf-om`Tol*M50vI;zX%9}`jfD}mUxUumL9a@^eOyeCbhHvai?M2MDYbh@ z=iLGU*SkWmnbxZi<$|_hRS=j8?l&Gy=?8G7%KYQF16)1I0OS~p1HyC~5>cUo$h zL|p`nEZgJpQONMltw4(p&z|7R!FI`$lJ&AY&l3eF-g3?Un0SxBB>?jTHpCpiPy?8@!RR5% zV80j4Wc2YQf#BR^u-y#fp05S9_?YlSHt24vTgWRxD9ySpA5CR#uP3r&~yJd5d;G+!7o zh18#E+E7?A3I;RlMZcD(=cK+ZX8V%E#|6GWTfV6rRr-x~p~Vh3VAn*H7+WXJpY>q zgVFg~o8-h<$Y_?3a|z1Gu7%$Nq5kah61S}~JJ=2$KMm)n4GahXDO>-;aRV*n0V$$n zVasD?-w!2LJ$no8dkxrN4@w}4S_U(~Bv`#^DM{O9>NTI>uU}COcuT4`U?vsz@d94f zCZiaIo*xx`h2M@6lR<3E?8}~Rg&BcY-hnvZQT$_@EuN zOn-U*Kr`>juQ9ss@NSs>ITjuX$eLIfaMM=pCHJ*u+^K+vj#FhZB0%K{xFtWYiVk87PW;nj#^dnS$&k7?^WpLpy{*mVu%a*FFl64 zd1D(j)~`$$W4k<+v_GGOJvWs#Y>kZO1VzriMS1ouc5&K!OYXcC@_%AR8!IfSKwe?C z*3qgL*gM^CgR-^BX4NNyP~VOi{(|oBG+u%(&ro;XRCgC3` zSY&qco9|HKI3X%C11NcCDcHwmUd@UdE&|8omIaTPE^5$OdN_ubNi8OTQG&uVWg#QE zbg`=}kTzB9Y!Dl`^8PRj^!TEe#0)?9TEMEnJpu5blhRF5nJqWU2aQFMQj$AAMd z=7BzvceN}3>X7DzL8awH2^w$2(OZMtu}`LnhVaSBBFT= zsXYQMH5UXd=iMOAO1P35G8<}Qsz4RBg{1RgS;+zmGx50 zVUmKOFDOjNjMRdJunrWMEaQp|hd$SmV6XZx-_r`cem} zUfAsg@8|uijbfSlx<2*Lhj2ZOPSTmYFBCpCUOxG5pe?piOm}(0En}3p2cU`KmW9aW zHJvxF^&?z=>NR&4$Hp57n64?-@2@Mc6fS+>J$fd5z)X8wKxefPzx1fC`bNl1f5l{Pl?a%*_6Ojp2)6SODpywwkeziz2} zR;a_UBkvC(5^&q0x$^_@on>KxbXssA+BKVJop*oixKSk_)pLRWm49#zR&h@sJU_d5 zEcNIb!`{+B zu`w^55AWK?YM`#%ae=w=h55G`RUq1iD4|}U4 zwJvK5Z5F55g(rf6xp%5W6laqqLo>x9`dhOdmM73@mgH~9`NI2te$IGDeF&n>YWJSt zannsC&731@TId1?(G?XO_Ii?|K!@Z1gcNoV2HQ3~plDFjm)suD)O)**NYtAwo^}&= zyRHBZGNuVcnKD6{-gSGG2zZbSi$qjJH|4NCG5Gi2Qw!^1xus~t&O&Xc(>3uCQz`nU z0HzEHqq>BFJHOee^;M|2YIKa+7>K_}8YmS&G!0@pfPGN2l&=qR@}z5{6&n()!M-S& z(r#M3e+_vP|5aNK7L`#VsK-SqOP%$>M3>qXHOj{Iul}By`msD z-hJY+e+Rb3w^zx?n*MgR%}h-Wq!md=-fGKgM&6Pz@;-`L+vZ{6&-Rk}ZAB%g=YRj< z!=x}B!oP;;ufY8L~}J_uoIu~MHr8P6Xi)*J7I@IQ7xs@eQ- zieetEpQTT2XsAfp$t%a4O$_v*(ilH#ua{K7+j1w>ug7cj)-Xgu zksme*22Chf01s(1If0SbsMK{-XM%3+t-pn772=*JmNxF$@gYoVr@^IWPJRElKF`i_ z)AQa#&}c&@r1W4Jz*aS*rRsYY+(}9^j@uhbAW$Jq*0gWj3gSetRk7p**Ppo#vAL5# z7VJeF>3$7{W`>lt@q|-L=n1&L(D#V#5~}}Uohi_x5kYXI0cWb)e!Y=J<*uhRjoG;M z#^ak9B`#*GDp6wu9NjBK-m|~Th6S)s>jCz5Kl2dIb&ymECX&bYzxuDa|J<1+&BGNo zi;`L1bSw(Jr^xDUqg-VI$CaZdE2iDa3M!ewi88H|u!F>suau2Wi@i*my&JeS5eFTQ zBI1l%4cW-b-jGN{I9+f4x#O?uP&U*RT^EwsX-OSOD;63+}274@DZk6~) z>#w(FB?8v2XS`0FKAOX*bZXq|S7N+F53MO@h!oqzkCCxB=?oF^>y@cB*&W4$(6rs zQ>M;nP-TCaSN`V;{2Lg99#xkGHX}C2`Ga+!7l094+~*@l^a;dsb6Dv%eaQLZjvGRGtPtDct6S9nnJM_l(HjPf`n}ujd z4XB=Jo&g{CKW$QpEMFP6Q~A|tHXV9EIBF}P`Eqk6s&78HC6d*rnUbxsAA#&{yEy{v zvYw#ZgRPtR&ujOWVHR#4PQ8|U{Z-p>>9Q>9+0HnL_(vKx2Zo%bt}j#LGsDE^(&VQD zzUk6f#PniMQal(aS2o+6phK?OOxpoOnFlyQ8BK2+d~A?&HLnJ~{sy`E|MGq`)Yb7H zKW2`4wELLnXmY4a)Iip7aW!S7vX|@zRB7mq)R_I+zWK8K=<5kXfN&hYPawp$$A6A+ z`F&d`KB22<+Kbo6y!!J(2`0VwEnA_;nAdCKvi1R#jU4>NlniL+lR})5u~5JXN@O<5 z!x=Sr68Sh9l{!wUpCG@%o1`gY_fz;pQ9tJGEPZf94DY|;vuuudj647fF(y0-{3IFR zCs91Ve+$X8^A)nJucX~t$LIXRRc`_@KrQn*I;IT$quNw=O)bIh+GzVuwZ8vLHX|Y4 z)j6%~IxHG9??i&iv5u}(cm2rw@vg6b8xcJ6;$qwPP*Lr*L^6dmZnFZ;@}pcISc8wS z?T!36BF)C(xG=aaL*&z{N}VjeY*?5koQYGIT82L3RY>I;s%9M}d^x+7;1oq*zQ2>& zpv`PpZ3fgr0sq6KkY{yyVyEm9VXJ{Dz)7eH$)8Q7l*yg?Hk(1?`<7+m+YZDN8w$6Z z?QI&SeDak2!UoKrwv0rG$u>;|OH;4*cbYDvTmYKJ8Y0TW`+F{yX%XlX8M*8tu-hEQBi*O= zFzhF4-0`k8NHeanTYHEh$s1FFkI_3`I|Lt%$hxii7!q9myJ!CAvuXuOuj0U7=g>Y6^ulD$|WSCzA2@|#sO-jw>D>=aN4EMy@mBgO}s z-MX#rsXY)glI=JdcU~}yyl$E!F z`i^UA*NN;FpO?3HLu$h_B?tiLu91CT!qoQCm4+RLjlbdYojxN;@x=a0?Zg@$(iZ$70Ww*BW*9LvevwOVi9^Vd1K3^c`*%gU~hh@pEuh0>8G zQh7qGX|t&dE9`x-?!C-ped9mu0l{qKZPLlVXya-`%HeYKwb~ROng#Q-2f_zr`(t7r zXPvq4db3AdIv3vf8PXFXLN$zGXx%@aO9TBsm13{|VIWYTOP4lfOf4EKOwl(Tt+wY2 zn5?}-;jtcpH<2U>yW2*?ub)qE0$ftJ=JZ-w{ky(0C#OlL)q~;5yd0)tN%-s&P7Z=# z3Cf80A|-m0b9&seI{_xT^t3WvO~cl!SZe%qWHejd_3e>6Sg&+K+ucXFa|layy2gkU^d}U;?Ni^EVDC2h0}{^Zvh!-i`au3_ zv~0-0l1c< z<5OPRqHtZ<(y^9kV!308e;!Nrh9ZvU4RR*sSxAWcN?;%6!tb|SEaeBglTOI)_H}Ivlhwzctv7s{jq1=D|AqoX zX6r*gg|k0>`g;3!Xx*;ZQugY(p|8_VQ{XeyEsN7uhKrT7Migys$;st^!DC_kJ< z0oc@di9B(V&pI&IzCLKm`&0!D|4TJoQTP*!5;S^}KOm(e2wg`RH1`URo_;^}t!^n8 zTW8K(95dr0_wF<;{@F^!I3 zFDg^2{ir;d(lV{{0!x~I1JL=0r|6IIhrYc}HBw|%lyeQIatDvHef3dN^&54R+$EHE zeoT#8ISw)~Q%l@F-OOSxxFhh^fH_0^#>&Dkm!AhF|HE{jG*Aois8HfgB4B}MYJs`u za+v_#R}yx3Y0t#_86B2N=vMZ^jRbImMlw3eWlm`RdRIG8%>T#C#Jsj~#J|lIqlZ>0 zH;)WGZlYkdoiJMZ0Ap<#C$v$g^vUtJ^TcJKzx&1vm{ZY|uU|YsyIPPIRNAZ_`EB)8 zqv)DMm@IngpBWDO&kWa}e`9&hfNK(eyy+sfC+etT0^Y_4&{ZpXZWT6Tiqw()T;|gQ zK1!pZzpdAu|BKCB){rXO+ZSJ{%$OFF{bUTpB^2h*HJ-gRnRK-)e5`NoR8v!7hX}Kz zAwVqs7w*5TUvx zG41J|Lj5?kNVNH-vBG%ByJ1>iy*K{&tRQ)#?GMvw7#m$Ia~2T!FmVEA&+lvmW*jJ0 zbb-FnTtxh|rtl2<0F&pDHh;&xPyW80EX<=KZhwb(X8IonSwy`+kzY|ZOkRKwgS1N# zmVWx)<)Yz2c1D4!7WC}6(C4l;^rhK6HLmK~fqsvVTK=Gj=igD9i2jVyC)X86n;Vo{ zspyAM_7tGDX|ehy$ZYKJqIKP>1iXg6@zDfYKATEt<8i5GfYm<$MAe^#7G}%Ujs`^@ z!DLr{&?CRMNNb0v)XU?u930SS8Zj#3L$3`XE;Sf`{_r1FipAOlIwSxzNTK|2-vsP4 zv2J-p<=Et#f_`{fenZ3AuMc>1-rX>@>pmAN^B8jrh`W5Xx1rP?F8l0tb($1jrGNzBM8g;f9LDfWhKiUbqzJoO#@K1+2!=R{KbCB zX$41i;w`n@v%3m~0xt~R?R&7*_?^>zC(xej*|O`LxBrg7b&2A-=jDM)xs~?jc@Y)O zs|G8V-y>8398>#Qmq?{?*Hzr4?X0NmI3L*lkl3K>bRkP_zsUQ)gFL%_sd&f^Sw)=i z?lr5+!4|uWhJ1$3CIr9^B#Zd}>wfgG&=p}rtHXq;}uKSV+*o=h3Dj>4u zeNX<5(tjxkL#{Y$*tLaPUmieqhTn{Ig-TBj_4}Tmh4`1Im1%AvP8&~8m)5KGJYm}C zd;e&hur3c%s%*Y>u?XnEN%7&n9-wLfXm#Y3dhpSz%fHmPOv`!l3<`(?RJC4`UWWIr zzgN|}TbG+kJP(0@#8ePLdG4tR{(1YFg{j}hWuw2Vrwr$jjF?Qcg^(eKy(mw}nc&8# zMZ>5_@ZY!p&uUQgJlz{COS}LNb;ln8+zA{)*g|rjcN-jZkHd*?(ACjqcTo#5%-2j2 zaGO5UTOT1p&^`D&%HOS}s)JsOO2i5(Tc57O>ahbIdK|Kk^M|}|gR9d5s2%F1`Uc0A zKI!A;F&8A)$_zAU!{OLYXJpq7abobfVTrAk1=kV5_o$R()kxf?7Vr)%}&bPRxn_eVP z@r|+Ki)w;V8k|XMAvSZtN3oYdp^H_a>o_|4l$`-BuBs`+F^OJ)yWQap`BmH)3_IpH z?h5U|?4x8AGQWGew+)z?n`9mOz)f>ntf7&~2HtThiTg|Yg){!2XbaovPA(Q|7kFvE zN@&$5U3?CJOg1`}aUX@JEDFE!vOJinop4*Xn*`m{*QQmLH04>Ry@b)&dn85&_kHG( zyO`R1PR6%)meb?x_l!scO}l#&^xS%|XS;msH(B+~hYIKpkyqk6s}4*(J1U!;Sh>)r zNftzd850p)Bb$mwu6b?7Tjz1n@j8POZhVmDglfu2_mnvL-qWetiW)r6gjq}_n|$kX zCCb0M?aM{8wr89fd4(*j`NBuvChn<`j0Tpf(U|*re9JJTNQ^W}h%}0b)Qe5lOHAI8 zNkCJuEEWDnhGyhL=2=h4+(g)Zp7gXh$C*M>?#;q1O`U!qwhEbeF<7;$JZ3^@@Q`p4 zUhlRq+fm2ikB8h^UeBh9dtvNuAbLhBLes^is4(#oZO=J3v(LMaR+tbcyK1oGU7uqa z=*klG?Yi#vuI~1p?j1h6J0iubS0PUhsb#!aB9E^MRH;hcZV8+p%x6hm`I!Y7nO+IGwyGEGgRH7X&N4P4SVN7x@XC4g=OcEoV}QD5 zvUg_I11Yj&T`c4UgdOW*zwtW%sb|} z@TM;<2pN#*# zyRt%kd*2$b%gr4#F09mgpWV_zQ_XK-Mm2Vb!0|^(zbtG+pyV|{N4T-nXFpKV`L2e# zRht4s$M1PZU_+a$8+L+jP=?*9W}wWbWEBakJB*(AZc;=9#9`fa%<#~327XvHOEwqH zL79L>F_V^7+DjbF)wpPb7ey0y5@yxQto9`pj8YIj6LRMb-W3Dyc|2%+C!0qA%2Z#b z^P(Ht1kMVH1sRPe{Oc*E+q9;w3Fwo2IgAubmK3b88@rD zG9Q~8=uS%q8+gXSS?ieOx)>EYP9FNMi35=}ASQv9?SbQpPD?9Y+7cOvrS47nI*;sQ zcB}QWDhM?2d5jf=Jvfsa4Nt9ky*s=@MH~KzZ>n#sI5a3<0YlE2)7g}hGod6+nU*4e z3rf4k@Z2-Jp@B@!NU8BoxBzQV^ZfQrudC1RQ-2Q^;vV+>+MwIabN+l-YcRF>>_Lz4 zgKLp+j4lyQj5&JrO_vXXvfcQD#H-&G$QXebsA*SA6;RII{bvc8X;_Y|hWK-48dYvE&} z#e~%K(n*UqN{Q6!tV07ue*#wgu;3LKJIs5O&eg@%hrlzzWPN%Cs&zVNR5o1dGs|k( zV2un70TMy?Blp9Ab~sZRNZtXupin#|g~>YaAFIn0oLYN(&{><^Pene>>juIKaKVR( zE}!EO(hFItxoIz6guaTTj6aPk?pt&m2Z3F4AhQJT3(SNMwd{@3AI#=I z%#N2*AbW3zP}0}3M&f?!^VcJSxHF=51~O*pEi6rV%)j3BVpous$WVVF{Q~}8md4~d zZOO2*nF$L|Nksg3hl1Uwc(`qk?PHVx%ngMy-bD0dw!bU#ma07<8F3!wnsqGmfF!N+ zI&L&EgkmTdaJ)gTvWN0QbRKh9bjHd&4_5P6OYn9u{?*wkV;LVy3pEeN1v{W{#>ahn z=93GW<4?*`Px(loVsmi$2JWoQGMH;>1ZXmEUEp9Y72^USAF1DvLJD7w3 z`^{ZqV@E=W7$xr^CLbkkN6o-z$bjzS7?Jd`8$N)tTCS1$-i|4*I5jF11b8D9uY~>0 zvXcFcU7`$tSdM;gm1D+d{UqFltrmlNIs(sK+ZW6mg4lHDmG@Eo|Ch1qNW+=F=k8S) z$H@2kporkbAJysSi^h2#snL;I5z%VIZe@JGe>?C0sKMX_d03C$&W#2(^TQoRH|iVY z-F;lg8{SovP%Ra$ak_>dbR>|y2Wddzpgcx*7=)?TeLG(?xbk zm3EU)8&|7rn#XLqhsNd|53h-Q@8y$8eDX>JVuHIh5tA&x`kev~LRk}*o=ba&A@wsd z_=5tj>J5o5m&iK|!O|ULh(k2HThZP}d0nO8VE5ELrTAd)v%Lr}vO<}rY=*s-0u|=8 z&xFCNgYX~Y@pg?Nd7&og5J{>@+TrUKe((C`o%fE8?W)&x*FfrCgUcFEq0Ty@Jncrt zRR1`UHAKQyMZ@olqJ-T;s=)n^`e)B=*@ld*8#?^nt6s= zUD)HLB5kwG;XQBhe12MWbP6VgF9s!7*YBZDo5T~CDvgpKlzR6sWQtH(R_{JC-`tq| zKG6)XeGhMX*NtjkB`P#X8m)6JOwdwovQxS}?qSx>JqS#Oi-JjmN4n6IP;myUSj-^;AdhR>m-kDTgdk)Bu*-v z=k=h_EPI*9r^H&l^D@2=T*ao$e0x;Ltt`#n;wk9U5 zG`=;dn?LL^ro}D1ZFee%g1lzX2^iCL160$1r zN9k&pQLdXzYKoFfx+r!Yd0jLC{`0(ZxNP|IC{J^wfxb}K3u=(C+JUXcLfAq!ab`rE z#fYiVSwjGXwqxtv@ilg1?{b+6NT+XzF~s4o>@awXCK4}VwO4jQVgRz9~zfN#66@=@&5d#LC;p?gH*0ZQkDJ1usmF1I9L!r*BtPb=C zgCT;CY*H5--bRj9hbjH`Z;Tj?on2G0?@w9FClP(DwjW@w70zaRQo#B+!428;7&$>U z$-T2&lpy6QFB&fppulggWpZUlLc)&tL5)?*s7o_p&u;^vNr=^V`Pn<^wdkYjoHDJ( zCz=g2nlofbc>ZutPr01sP;ap0Wxq^U@Noil@AwbQ}*aF9$dG zQx4EHsw>_o6PI1VYLtcJT3Xe}|ZVLG~ zFFM`y5Zc3TqB?pk0&Sjm<+%UlXwLxl%ke9h-9myZMMaaz=bN!o}+9XGD3lB#m&N~dMp`2{4X>@@8l zJuh^Oy~*#CH~Znpdv#6F@n?UHUwT|nnBi#7qA{MPgrfEl%{lkL5N{^@wv62?B(nTy z6@svvDvchpImjxW)*|0Q*K5U(l2;Sv0OlZn?I3&1jeGfJTbofS4_PPT^`YjR-x}Ov z4M*Roeah^BCW{-!RwOG}rRr|oQ-ld8*V@dp$^8XC88_WAh4cmsaM|sITT#eQUiXir zD+;qDu442+b>);WQA)AS4ElDJo<+qSh9^zhHEH+xnZgdSsy0{c=BEPO9q)yOVTCQ+ zlDwr_h_OWo?UWUGkOrvLP6M5IR>vG@5l=L}F|FP+t|Q@TG5P zDZ*uhbVxel2OtBe{_9iSd?oUt$_j}(xH@ecejXu zbax2|NDE4r0)phwFmy-^-8H~ScXxNk&@e-t*KhA{pM9?DJMV8UX6AkFTx;F;Q_>u{ z8v5?7AO5rHA>0jB6vmkM2eyhSfTPvVot-#tZ_P1&2;0iiZJJ~`#3SV4FZom zF}8rApoXELMj-Iuz%^vDWE-wWM9#igVD5e}LbNV04DS(eeFkon!8j|JY>C=qITsmGC>cb{)>7E})+ z+vK#!?>hQ!=;bMNez?*-J~r1fw$QmS(AzeU)Uz-!TU==_qmL(?EIpH?qyIAW69f|c zJ-f-u74)LnCkOX`OwCVzr(zJ3K$oIQ%0+*eUYJME&*QSk zDRA^g3!qbA_M2y2<-#tLdB9XS&|nkun8g)`CMeKZUf6kizi3&DrU4C2X@m9>7m`aT zBCUP5?F?zssRV>L`M0o9wBHtB^_rq>%J(Il7H6Ncn)LB5-;eS8`i*eM;@8K}t3#m9GTolB>cOGtMA zTHX^+SX=#??cABS+Lhn=Q@3OcIGgS=aN5`0zI(i-qravS99~(U6Au1SZ)HfB^juK^ zoEjiY6=PG{spXqPWiqx)^u-33pC*^q)I?J4;#r3RRdAwUbot0z9En1Z)gwJ`Su;8B zk7JgLiUTAA`%PT+DT3|w&^f5v`TnG3qU*h?(QVgPIDf8usf0_(Lcgt$T!5*CGETWo zkCaygz)zD1nYmweZEx@T>DtEkL0uG+K1JhVp#K(+-=6*l z(!&;3xYhAIg7GV85^8iUdBHp`;}+L8)u`O(w2Iw4kUWZD@XPc4MU!#tJa7y&$d1$- zapkwfr6oRy?s2RCekZ0ZBq6FAzGl7DHQ8KPDxl?=n=c&q&eyKt%Y12Pg!6=hN;xvs zB=}cNiYBh(}Z{ENEKGWjl3IW++^r_04^iy7UEmuuo~VdAHI8Rd7P(ofxNUbXN9 z_t?0-N$|;ro|Tp^eVS^f2FPZoL(Vcv`t|*gtqL~W z+Y_@hU3&S@yQYXD3RjSl)@ilekb8(C&GmTurnRt;2cEmx;9~nZm%pb zN^-b(MXwI%-ArsvSb1xvT@ED{IL@~k;vn?olRkgO2LF2g6{Wi2+I`13#5JbmgNeQ%cn~UgQp>PD^~Ews{Q0@ z7jy{3Y4#KaA$h~K{Z?vkm2knha5W=>AQL8*RBs2G!uAEvCsHv+V-Hi65B;HfehR#N%s@A*JrJmYgcz7!xMhF%fxvM(OkaTSE_bQi;i8U{(o>*UAYkP(mk$?0!Jvr(Fpw~u$yR#0) zT>2s9`@8+s1;%_SM z!=BpG=B+0m*$mWZY{c7UXy4c+$!-UHD?hlOeqTIDt{3TOp*xA#SYoF-=qSMx&Drbq^a5a(bVK)+ zYP>Sb$A5>BG1trcHg$R_=du{F!%g z-sl)te8z0)DwfWk@a!5PL@yKOt6B#^I>#cBRYK}K0nOSPimLtB3D7STmZS~WKG&%Q zRSk|1CO}?ZWCsUrF9+XVC0fUomfxmpt-;?(&V2Uv>HBAN{?hcYC= z_k-0%;5mCiXJw0zRGif>U`GPgWptnK6+_nhSoX8VjBoGv7hE;ebz%HR;!EevfTTh$ zdTb2O1FqS@cCqM5MB>sB$j}t?;U&p#es+ZOczJC(Q6!rjJG#qGWaAc1@9T$8=lZ$l zmn^-fP6a3h1-Cv@+Z;Q?y?1{gF=K&Q=>I;&-?#wJ6v}9LdVX5FE2*2_tAQsEY~K`b zP*^yv;@el4jkJ=lqXEfBJw^(U&I&?B=nEtuGZpxc5P>|g4&M4MuW9dBg7^8axy9JF zp!~&bHwGRc9QFD~I1y$dnV$fWIW$dU_r;)tteT$WO?8%|yq|0dwv4KUwtQx{m&5II z&lZkn25E8%sRqr_+Ho^~7`C3bDw{J=)j0Y6`8j_-8Fnhb*!3G!jj4=5_iV95M7aX* z?TR^5<%;oeIqKr$E)!A>D%}3`jbumH7&(YO@tfDi=|US9qtP)_X{gZTVo6DRDdy%j5EU_7+gT%0NXbkLP>De!SlzU`kyJ}Rm`YJ`vK<{|W z&_ror@;&K)xh>#V=>7cz<}wC-ip2nJavVfl)B~pAGT+;h+Ga5~llaczrFMrcbDd@_ z&9^trqSd|WX+(x%QMbr9(oHTX_ON@W@l{0{?cZysPX`RjPmfU-RFMR-rZr#|GCoU8 zKTZtv)4Bi$him-=NgIN>Oz9st%iVi@c$FBX=|hZ*33`H7nlxSAIxKvcB3|;Y88ys1 zm=ZFFs|@nCYmMg-S-M~kDw*47z5nfufy}A(1`JFDlMOHF<{o;fmG9<#;fqLPAmQYX z3$P=-WfRxRBL=>yi~v&Y9uX%1`K)Dft&6>+mMwmAKc+)%ai zJ+hdUu6G$^P9I7ZXwXL{n=`3uO?SUBQ9m+0yj~Jz8nDjCIwhy* zFUkyBT^w)7pmrUwOO*|I#ph{0nkQRmoNCyk*;kXgBHRBy8vh7#Ox+YR*7l*9cO8x; zI{xKUY0$Vn09c3qNOrt~bYTUVAS>qpAXrJ|X=xDut-&seXfBPYt;%<5>fh__ZmPW) z{u#G@@L=|VIl$t7eKT>sTa~c;_$5W31|%3_eANuEi~7aqIkvmA=7~w{|C%8oOB&O; zPyU7K4$njEZjgMOleYC;6SZ!?bTnW};+}mkVh6a&U49BAK!i=_(tmx`?9VZ{6l)W- zI7E z%^~C)~EcMYFw2i)mDDZ;It>!LS)m_CPI;5VFztJCc zb-cYxv1&4c{hVZd@H!Gk$=jQ%jB!XdSg`0JGxDWfBeu=)_&6LIUn*=r7xSws@aLD8 zHHTshA$o~Pi%n(E?&2P-Qz#qiwl=pzh(o}-x>AO27dg;J^f^Z%c(fjAd6xt3V$8jKqaPadG0@W0bc z(}B!ibHIC)v5dkE341Ln19AxqUGDCjuTkJi5M4lk#urMvtjbLIB454x>k`f*MDkch zqwe=j-LE=iunafx9<>-`^FD0jfbAcen9}!!tMsv zqr~u#R)K}3JIU9+`WhShOYvj^$!P}^zVCi-jre~ly3f7UHLg@Xinw-=W}j{6;yAk~ z^LFEXbJ+yENb!P`_7lq}!+FKeFD%d>cQ`L1@U#e5$$%Z>WiOc4)0HFrD z#90-C5ORbtdE*3G6JQgD1m2uoi`l54EEigZ!XU(CFifnIN5tFbIFE_H^%so(^Gx#t8d#K zK_;lu1ujLNK$A$O&gg)3E$(BSeeLOswhmX;oVhS&C+DnorIzg-9pOffVk|tIce(gA z4uIHEt0ANZ0~LpMYOwi&Yc}j|@#tBunZYZF2h!G68=4KOaF%f($-G}oYucdK%2Qf2 zg#JILCH41dRVaBItLHu?N4B+&8nN(_K^ZQBVWHS{|WnG zvDDk@wy<;!oS^TkSp77)so9A;$Kle&&>dQQ6g;%zO2H_x#9{m!qUJS z)Epy9n44E{vsmAvw9^G>OjannHhf?4S|``Td8dR`J)7-6IpM#4#`f1MJCp1iPYfOM z3{QOj0WACQ*mm=S(A~!IQ z6R|yCL4W?vN~M?k(7u&JjSA<+$)JXR4WmYJAz4e^E_AR*q)+f1> z<=RgFKD$qk|5QT54>iXj@VL83t~jyH%{7p%h$k+}cau73rTIyJKDn2Ba6&;4$HjSRl@A6dQgPIKEfB>) zJk@NZk;T&CcQ)1%}c0=t+7Yb zA%{Jy><;FNrP>oyV8K}A59J?R1i<(5coT0f<8oQE_Idc|C$F(a?xAAwOI}}0_1r<- zOjdO~cr19%x^Fz?|6hWfHEVfSUIoQMJkusimV6ibCx5&~gS;#AGP&=in&tpLBsQeA z?pHgoUrJge>f?OKPwySfhPJxueZ}Irvinz^ll=S%Hx+$wW4+u-=5_dqnXxFSfF-Py zIR+Bzw6OKHET5;6&zS=5=AKF->QYC7x?+o_2R^&)9am=`cd$)Z$wFKVEx5YwG#o{4sgIX|(^Ql{Bag-CR+ya3Woly}xu#nlQa)VhN;6lZ7Zp%WQ zhqmTc9xKzvVF!T!7pCO%*N+SciD(LrMO)_v@cH(n4~?Ap0`Xs@!E!K}Xrw{;(<`LP zc~bUBX#uOE8$x9FJajvxb-LDlWJ56OEZs&9_f@)qCVqOAGxGkc>Bv$hv9w&kTV?HJ z%}~_Nry%In-fAO#i-VKdwttnQ{e45D2RW~&Q?1xS@u8s(9Z(ZX}UhI#X;)RooR%;p(*D{xM^@r zJ8T^49FOZsTQm^&beej??J3%bJA9=OmwVXYV}Xc^>6064(@U;3RK>1k^*Fh|_4U}K zYxE@2-V<9{8;Lz}bssq!vA=I_VBqBu7rYuhrd989S4#&9aDv}X=K&#dp1Vzm0!Iy^Y9?0mZWCqFGO&ArU=!Qv2bT4K{WKmntWBa z<-5Ic!#;?xNe!m08miB>(C29pb`~4*#dfqHmGzNQ=DTU*Nl`w9IBoZ@IBVU(ufFUc zc8GRsIk^ATwh;UsA8Yiy5vN=em%?!n*`0!6Zpfh6V>Hss1M~Y6-Pps@YsyGp4qn=d zFEEzLG47hOFVkN8&48VA7t}Ei4cP5>JW3Ps(jzh74f3lyP{j)i59fRipBpQo5mmT$ zGX(D+c09n~<^;)yQ|zdw5awQ&Vs5v^=LRX0&ZYXVDqYaI_p~&0O^F?xpXTRhsJ<2J zSUZ_FC`ut?-G-avcc|b?5%=Fa36RQjtFDY;pZ68_%NxCuK?U*hv;w<+OCoYb8xk+~ ztq)hMtxf1QPV3S9#!6S6>Tj=w4t}O4rQ?ofT)AFCX1NGxQShb0U+yA-*5Fc*jB%pl zcqF;YNkC`m*yoxX_Cb7ERJrgPl-zMRZ?!vHcPnJRJ0oHt`p-%#Ac9llj@y?@bI(1*=>r{(`zs-foxnvL5CzpPc?+8+` z%2x1>A!)Mcd1GeAClRR@{6YUL;ORhXy=DixjO+Ez-WN1K>}T$t8)Nz@DFyu@XUU5g z632&zEQO)>2ES^pUg4kgUnd)+vVKlkR>PN#KJy$IQ4q}eobpzHCS;kowVZ>%VhIb# z;(CU-z9+xEY?$#dzVzH3aG$EZUEJvKV))6bgavi?G2af1vS%X{2pnb4TT1d@1`N}3 z9=(X!X65-kE22_?7Gm<)z<7(xhR!D8-Q=$ycINjJ5JuT6Mi;Ii|MQpp@V?&hSSEOh zs9<0q_v2b07IlXNz3Ew-K47KLzbdGz7k?!gu{g;2tf)H70o-ugl59ot(N?kzr>)WY zeq(I*H5vN^Ejm3`|J0b}4AfP*F?*csw;la)O~qDSy%!tdO96v9Pz~`$Ww+IucUU;5 z-|PB*8V&g6lm3)$AQ9N9#t6AuMA0zfA8jfeXlmjv?N^^abMGC%v{%x?@4++@&Z@3a z&QZ{?kSQ`0`aLmHkep{&YJX9~XMyGygR-+>0R?=w-k5Mvq>|-)5_yRzXBxPfN|=c< z3-lx*&hjEJ=XR+;PE6*a+A=7683WRv{SZ~j#$r}lV(AaWn;&dEZqH5#NtS%2UR9Yh zF}{{9aKP&F%cFEKdYAAkXd@40%$H3G!#fUd;-MC|y@Io2YfXMInji9cY8`vFk?Vun zcu_`8L$5m)^_WpjRhVU{LO+THpQHBa;6{P?|5MHU9qg=G)cHvpY0Y~zR$N)sT}2^B zp&}oK>%XLvyYw^V3c(J8-rTil;CCKGDee_H?C*})TsKWMJgz&ig`QR{t#{xDml)2% zO5%K(Gt}q^e+8WF?X<5f&!T4^vS|%ba9yLqdV>vir-r)GGAW4A-7D$bH3|9I)tN(A zqjB>n(WnR=mG>{zVK+r5wn&f zr_CC2+wo@UciRN=0$61A*PD&ClwnI619eath{Bo` zXg`lTc$DsrY7s3fUYq%A^pCMMgsugAQ%*&?OGvs?-PRTJ8eiAF?@ksRDXhqM2N(25 zB=TjFg#KuN>R2vfjf(n&LYP995(E@smEY^Z>+5$2R;@DliRlz8iS38qbh3w=-&W2L z8u%T5>H2hsx-`N|!Vq39sw}uA&$63GiGZk6_3}yha=SjmB-k ziODaeh93&B!=RBuvDSFnSG`LGK)n{&ktzWaz&u$IHKmRkkR zYqqWg#!c~TEq;jNTNber_S~UQ%aMJyLO3kmm(tkq$-f(P=D(VS?VxxFRK4J4E#j3k zUs-vyZD`W;EL9G6_Yy^%+TH5N=AzgnjcmR1LFHNt$)Bp6$Gkv1FmsIl>r|>nfuf>S zD`(3|PoISWTdwy}My-tm`sqRe9K>LjrzVZz^WqrP>m zv8IhF*IL`L0$(+&HjJ0!15`*iXHl>88s*@NX-!yOx@43_d+taMc*Q3EIjL~#-tUVY zDM%*46Gwt4~>2w_22(rSQ=_M_3e zX_ci>x}xeF^pK*XlrH@f^|y!A7C{ zY#V<6=6ui2TCof4!(SE}^u#aYm4c@n^?Nx6{KRD0h13D*=PsZ{uFc`W8LU#2PfLXb z0%DAg9`M5m*3nhTNWJM`y=x1rp`QOhm+4Hn%`UUdAKgl7T2XEAI0zDoCB_NweJlL z&4{6YHe>5Iv{tG!J*2^uu~g#guvG>99@FJw zRyhKotggs_O>=8!*_$0n3>*}KXHP6&uz$LTNx8+4RJs}Ynh|cU*=fcS7F?qUQ9N5+ zeM|S1Z}~GLmk8jKK}4BoovX_tY@GS&Yn_^HevhP!^h7?jzb3Lo<&L)qH-#jIp^@sU zkDibkEC5}&cyrCg=FCQEx(u&(3ZsC7!`Cu`%N>(sZ}I#ck*TLSDa|nnAoiDhzeoB8 zDlD~=e^h!or0O4+ye`qr(No2DJxu8E_M9HQ4C?4)2+49_@l!zy%0LtOu=CmuEC<_N z+<3tJh4B`JJfU@mJsRJem5?yHJEPo{t!U8qm6|T`)EkP^iOEjJMxg&wB%bI<{DDQb z{U=7wLdf7DT(23+FyGOuw8$mcTpxfOyxYoT_k))=X@+=%nz`0w5Z zI4Iow8;0UdMaJyH_|z<{{Z#&)Apu^b^c|+(Rl2uMA?Vbii<=#Hh}n~OQd(D8g5KEp zmtVO(6Q-PW=S=0esT-SsHGQPRbGg+-J`6uYXSfD`e;dg8nK)oJNO5W@B#oK!dYUvP ze3=xb%khz=78PBiNYO<%-4yfeHJoyL&U57oH!5=dvjC@lCy$+~Sd6dutHQ zyZLSA@5#v&Fr?%96ru5MeTV&O8FWTUsA3f}&S3%LrZt3#&ciZrV!sObY`_~vx?1iio(ZI2t5OY)3YbA2NmfsK7 zj;{)DgZK|Q)RATsXxt=Kj{bA(XUi=^rN&SQ^xUB{%rv z8F!Dkab?fS<_C#g_qBQqG?&ep<|ZZ`Nl~>FbTprBjK}w|QCzxKBvj0hk|)fja^%b$ zWiOR*u=XpdA4U^<&UUtYzr+NvgpMlBT%$giYLWdn@c*FyhI%!aQA^&epWgOr2(@%k zgbWn98Cb6pr=RSjEi68IgpfI5W9Q>@BMtJm?j0}L;{F3Y&RMajE9PiSHoaNkIzJE1 z4`&W?;roDorO`dRQvGZuyps9dBen=%7M24$MRs!IiVs>8)YR>RgDe}9ft^!8WSAr5 z&G};4-mB(x*Zp?EzcwB)n4RMDn)tUNt4~-c28)aePf3N~7^ZFf5nLP}jv3yPi)A*8 zaPHYJ7xgTlkxrzg{HUif9ov?fnu;?I%~Y9j3D3rZh~U`dkLRI@0u2Qe6x;C;lPa25%7r1bnIt}j0E_S{}x}oIuAtVhfO2# z;>4Lv$t*WUi=Yc01*nELmj!}N*RhRjW5lnU`as1F`|A2VL9K4E9iRbH)h61ry!8K? zo}p;&=(W6SayB4?ozm7qW4dy&4LR>uyuu$4>~8DIMw3o&1E3|{OFl?%?IiQY5V=E5 zC3g;dWgK1fRc*Z{rwzMIkOS$6Y`VW`z>j*jmLl@&ngK5_JG+wsXX>Y(6#(bgZkM7W zzjTzxKXspepY1?JiJ>kDn5%bVQQ{^sy(sTk1~(j4Q24zHeh$Si`1+B=VRbsYe~}v9 zrWhtCC(ej}IMwmW>7xsg;t{eo)?A>R_)n?*-_pC-|392CkRuzk)pJ8w4*qnA;_75h zyHV|}Nm)N}B%4C?^_~%&u!qQjb(VS+x_k{68-Knq;O(AH9LUK(Txcr!wK(y2ewnci z{@Lj5V_w-%ZbT_xUq0Al6tWf&_v7k@9Yb>?7zlJe+(=FWVhJ^T!Cx>lD~opIhW~TS za9Z4voVC6x)R7Xb`H8VWpCktddZ|gG8|Hupt{MM6n|{S7UTn@M-;8Z&x3B%CBc0nY zpNG}s<4_10-CC2}=H^(L6yBoB(JDNai5~wZ#s*~d5-W4~P|YbHyUWPF(b)mPvu^sS zDc<-NpwsQTL)hnD_D{-F{8bt?#`L}qgM999{1vW(A+ zP#u868m-6iS!JwbpcIzI+UZ>R(9dD|N(<>bG9{R9 zr%-5Wz3?IAE;-?))qF{FHqgKh|FMk$E?&5c>?y@RgJM9TiSt)q2Jm?5dyg&*36kdS zesXDAoA_sb(q@PD3n$Ntx4bw>2fVqFNAzYTjx3DvyVuiM0EY+dwF7OF=_%pRd+cLJcYTRQ|I6s znU1$({ct?aGz$w03wU(+{eZg~6~Y!F{^c_;J+N`&R(OM2>lKOASx19oC&sFNq_D}j+CFLJPtr&%Se zt{`N!^|8zs>{uE`wgejsrrlkuow>fwg_+iFG+Ol0wuwf^4a+#Lau)8!1kmFkf4N_) z)X*(H@^G;p`L2YrY0_M1Zod1NS5=IzcpQW!>KlyRF)5}AMGD=Hdw0{c+T-}ab8{vQ zJapf(Qs&A^mOg<45T>HI4KidD{&Y0`wR)~oJzG=Yf!B=$0MQaJaPjV5+C~L~YZRIu z|GTZ+`>U;dub+FU1Ti`ua}a?bT-nvzl*EBeyzi!M`7QgZ^Z|1C_Hxr)X?@oGF=s4U zozrKR<~XVU+crC+gv!q(o4gX=SX2nPl_9ttBi z+FSncu{0AY6N8V#I-;WQSO^2`i8T#R9KRb>yF6KrjbtN_r;%z9l$Ym9jPnuqk>zf? z>OQPr3xvvdyv3)*#$-9_iJp6*`K(XF_^iT=SPo8$PrI~zE3-`I? z&V9Rz0Yk>q3&%!J6X&&qf+-ew7!@h<(@4vjO|O^_oQGlEUD;3Oxs$o$=K=ESUolNm z$kaER@_^vL!qX7Ngz$tnD#W%(VMqLffR$}~5$pcg#PMOyw>9x0%S!ZkP7lPX<7h_g z62e=N-rIugLr7os<;%K1-~s2nX;K2@C1z%xnwwYJ@Tl0aRu8q*{|q&}jOs%tzX+FV z6&4a8FRocVe)GTg8OLBDPl2umRdTxa%g@|u9a_9!vGZJoa9&PTR94^iPeYwWNqloUs{46pBjmSN~R3R zynE`E{h)3F*fcBqAEMJqv%E2oKhBe;_!#g2BJS>B`KFw@(cSA|4}N4~Tn>ZKr?0ve zZex`NED>099P{1XiFd=L=|&~Dd-=IEE7M})r1TuI=gVB&2?<~hKTQRuD(7U)ZK+P5Vy0lJxcUTQUoW`kS$$j5`w3%KIUX+;En5$up6;LI93GM z&SKfd^y5b|?UPFS4hcE2{P&dqS6A@h)qluyvSBQC^&uL4VGi?xeQtvCEcN ze0g`}cca1PP&`Zf@YFPSeih;17;=({4s;;F#W<8uOi!UpsBm*1hjd6jj>j|-D4Yrl zUdUD6Zg9$kPpDInw1>31xwn6OEE({NVM(k+v%_}~#ZIIg@BzOmEbge2#9PHD^8*jt z=ISCXGC>p+d{G@Hn<^16550HVN1K8x8x8DsLY-X8h1`E*tq`l>^O{C>6>h1Xc`{Z| zgCtTG+=1?J#68D55`s{*Z;v)O3!$`2{vxj8q@?+tHZ884V`XjJ^i#L~ z^8-JCdP_Rx08~cu<_VzdJ=eo9ll3Ze(XYK47Vx?6-lVGd+3#uKvK*B!P790s$F{w! zc{I}l?ST^G()gAWpIX*Bvh;c(w6GG6N;bo7`Lnbz3MW(^hx%##p;BV&epLvbx95y@ zT<@y{#KdAa)N3YH3sIJ`%mTH;RnXZZ()9FfZ8>^^*O#$ z6BPv@rx^j)s09|_P6t*CWjRgl^u?ciTr&+Ue@cw?KHtjLN8B}oiJHZ0ArG^FE6>#{b*N0Dt-X&@Zp;$ zbS`JXf{^-_!)_>pO74?7`{aHuivw^k=ThFuDkgTh-Ig(Y$E=4l9 znS)J)slwwRJQ`F=9oyIeX+hF}0qM`rjRgv8<+%L;vf*xd&|R8kUmZN#+Ash}?tWj_ z4&q}V7((PaK2A3McJaZOO6WFe9-Z{BfLPC`xkP-GI$+nPzaV$~w{IhA1b{ z1#U`@|JhRUAkLA#EQ=$8obtT4UFb)n2>r?U&}k}JKMwwu>eaz1^Q(JI8AaK*(Oz^C zDgsJHIy_|@tot+m8BM|5j$u2^T`w+w!FI-z1WUD*u3U4wgp?A&@4sFad7n~w-1U?9 z50;_``ot;^W)C=rh}&#>{0M7q>b1miXKaz$F*zIAF5H{y``%22Paoh_@{T4;oD#W# zB73o8pX&6*(BA3uC7YIj#BeQgc4q@Hq+7udsm4$kQahfmKV`kml#;j;(}33iT3w)b z`cRQ{;YUZy^`(p3#K&C@Z+^C$hIh&!Mb0{{Nf*w9v>hN$R$f|F!?%0#%NG-W?7i5I zg$yn7{Le56O@cuboXcVi1goodnxmI+D9T0ign1%(wjDlSLY* zWouDmLC>Y=_Yl82YR^qMss6?{)Ml{*o6~b(rB?^;Ux((i#yOC!FbCLG0;4QevUnvZ zlk;%HKEcB?`w*|$6QEcDo4G!(BI9kRy@fVYd#}eP@3HpN&Q@ze-F4SX9|S_ z^@nPN*-@1Z{mN#jD18q`C1%Y%cQlAP(vdSm8ogVNkh^C4IH4G^WI(j!>Bak|JMX<8 zDR=*=qWxDz$JYAnyUBP{(!6c|;rG#5K@WS1(BvKK`~3|D))+!U0~)WjAIGy1A1&RR zy#e7_%;;b$hDuYS1%Fl{T}wpC1%2ST!_q1j|(1WsW+Sz{{ z&yDh8qr=5kE-ALN1Q$%=uS_YiesSTDP~H6a{?X;Z#@CIr6QLhBk|y(;-%Lnfv}PX} zUk++E&Fe*zi`BtUm9Py(?cl}+3wco92ySMMFNt<4=#S_T22O`nJ{?vjy9r5#8ZnY7 zq%2@XHJtXo0VcPFpCm9=-_L@}D{1H7F+q<3&29e%#jvJazB%zH~E9^y6w{pd$ykK z&|{+QjMN6iSVf6*kvXn!lLCh|*(Zj*a8;&M@Ym0yUJ``*QWi%OmfL(}6n-Tw?) z@lQ!U09_eo85UMbFqW!wEZ9j&u%x=WbFzP-+Not$!Wi(bpxshcx$w^b$p2?*Z^iK2 zeM+se&-4MN7Rl%m64Mk*B};0VOOi!w-~YC#pO$h@Mzp^C&f93;_%L6uG*X+OLz}SU zc7uBk&07`GL#f{7y7mX7l)PA~=yw2eQl|rR_;b$SM4~#UWmgu^9zlhS_SCZvb5T( zDEg@?DWYSm@fw?SMpQyZMx&9k+(rE4P`81Cz1RJUf1N zeQg;VVg1v)rh`ffyENf!ipr0ivWZFWi;5I-)XHn28;v`El24P5#Y|<8L|jsRa(FeB zXo37_+av84Nt`2FD{;F~^SVn@|1iqebSEv+nsC!|KhW$UIuRdT`edirA9l;8U2o^rXm_<9$g z57ikeK8}ECze*mUzS-u6Fe#zw^4kuMmkb^H-WX+zi+LDmCV24$B??Y#sZ#C*mKp>guWD6`%S)V?9jg;*t&z$3bbq!LO}}O%two?2Bfo)z=ck3W)M!z6FS3N=FADY0@u-(IB$FDvppp?T_YbcN zsx`mrjDnwWAS@wJPpbjshNMAC5sDXEWe}CDDe?ELfJ%QvAc6lF-1bHVdusOA zG=Q1JlI#J@Za5Z6`swTZgL4bO1#Y&m_(89BzpuPYKfFLFOMXp1T{s0sHd;HV4!K}P z6N)tfH=X0mA6?mtqy_qIm6Hl%HpayLPD0p7dyDJ<7f?0>GcBLef%c9wWr6T|_5)L~ z7gnRLN?fZgHlQF!5u7B@yJmRTMX8}3f3z$0Rl<5ZlwN;(JtqJ^FNiT+Dt_x7K~4}5 zU2e`c8}_f0cL@MoEz=V>_h3orcC=I;e$%`=k8@w-`MqBZdiL85JguiK2P8ea`|XwQMc(?gLt?@A_WTQe zww#P@=l9;^kB4Lv6-hq7x7qH`ndNdm9q|ai<{S#UV$@8#r8#IDosQtN`!TuZAgb7g4=^XHTgWp>cehUz~z(j zC(4rcx6MqHrqN4hi`1yldZ`N545|NEzLsOQN;|2PC5(-wop}>pVaH` z54aBOv*pCx2=3T1QI#-zN7bys#Kn(i#t>+N68d9CH&k*moW<+aY@Ma0nZ`yiplQRC zcE4Aqr{?nmk*wSTR^Zw2a$%UC{`4@nxygVbnsj`~2*ILuDYzsWZeF0$ohn}GDTOm8 zZY9?X3T4`xR6S5=dh4=vIf1sd7F{>^A|B*m%fcw=%u<+zAhs9v9aPwG48l5dySAf6 z0z5>Uu%+UwwU;jiR@t?m-m!!O6O#6~v08Iu_x&8eB$8zTk&=o(s)tqIxXZ@IWYTB= zs`eu)zHlQR==FY&;il3-gQ-vYz(Z<>G((Ad=aE(KPbzePNY7BTdVjHF;>=DPHaofd zb#9W};Uoa4tM@G9OOD}t{7rlP!AA0%7tgRrlfBFqggr#69M>sQ9GKL{tS#+hD}`^B zFMZmo&3AzXw9kcizyfajiY0F9WB4wrz-6~iw+ib=wzWvG>E^h(pZ@-Oar3OkVN?CSbSk=Gylj(^pMZe4u&QoVYb znEGnQHmprT5G}&AJuj6(7Qis62=tS{0@I0Z4ey&4wiDf;U~5w@>z}7)7I^KYf1^u=?AxG3tE~P*acC>C6|x|0oqowN4!PI4uVZ6rj5R)! zYK=V4PA-w2y90nMD_;~3cR>I~=x$M>cxLJh1eh_M+!MsSY&I1CS77YNfF`#_+G`&b zpcDrP!Hpp*Sc$S7i+a65@4SJnnT^kFttNY4L`uKUjMwwpG`N0l0WUwo+r&Wg*fsgv zwFCMsWxRQyGU)yl-JoM}nL;-ajtgOX)!)%xq(J>xDdPRjTc{5EM*zi)gw6~F{HD!C@c zTTZuPL*)KFV5Rq;@ujbC;+%I;&BeRUYXn4Dqy<}*~bIof*o=Rex`$x_@vG13Xa3}W`tU~4%D>l9n1d6 z#Gr!*yfTPrQ*04Xn%uLAU$oynA>}i%VEVZzQjzM-dP?}r-40TetDvP#XR}qP7fQu! zI+1n9LO9x#EIGNLIE%W|2B=E)4#bBwnBJ+*_XpInTDM*|4&{bVGOh#7V|fc+ZUmeS zlz8E+-!MuP?PbBYN z!x6KUk}+g&_oJ0PooIulPmU=(XRH3t48W&8jj(~u)A}}40RZ`#^rn6<3}j`Z8k1i> zmQ~4lGs#4s>0!J}YPobd69X7?X5)|;&$a$>AT*ty^^=NXx=j#R5 zGbHr|v_6xBYJRW1JmlxiEom^mbA1=;$J>VTr?>fm%aQ2wD%AU?-|=)%E}2upc)=A zC1dJaY0q_TmVOk0h`G6t7k51V&G{DW4k}y&Hj>Ph7CR)Z8o+vTp>tNz@VniZJsYxv zcU0vc+IlmqXR8%RjE&OXFRiau4fl23UfQMleCP@T0)5=Cpq;cpzB4eOeO!c=Z6+9g zd9j~}xE|@K2h|Sq>lOFT?zxpP`}6Xp7S|pt?&Dr|5!`cde$BL-ut?)eeL%NHT9f$k=;Tj79W`1S4I`( zCy9D13T4yeFo-__MSY8|?|VQQ8UDy?AZ)z?6I?} zAA!&K%c`mzj!s=DN4b#f)YDB7}a z^?nKrD8_OGntpUDjIS(l-5>-nhq{}&ybPUx)=v#34`jjK9p^LUm;L@#3Ju zyv7IoVM+I8u!OPg&!!?fW1P`SDu~X(0I@dqQ1f0n$K1rq)Wn#Pm-{i0ix=iUa~kQ5 zKzzs;>HEEf#K{EJWjz zDeca$Cc+Sb`E;OXW1rdnY1U@lcL8whRE8QE|+Hk_$0vZmoNdcs6$uA5O=BKpXGy+4U ze#y9x%}z}(>afc5pEyO`DrNIdbFe>*o7iZ9mIW`7?+`ILxrMh!WABsaY`Kks=CRJQpk}Br=Y)xHhfj3=eDc8tE4txtmx_2u23avE;1X0u*T} z-F2}zCA$(@LFR#ctkd$tkgj$mdVPwLL%H z4)@;bs}g#!42*TZ3pNf8*ZA)G@FP@f&QP@87S5rJ-p4xZ$(O;iui5j9Am*k;K+lPv zP%6G*b`oe9V$YtMTj+Q#?)W5!0GC6*Wx*O~Jvbtms5h@^o#s$}MI2Z_7khBETjuU# z$H`PEyw_pD9vf?38lKchZQ%=Z?z_mdIbjDHl0A7nM&!tPzHnW(6ab{>w37iA(oVuW zEV8U^^9=R9X)4#{9Pt0%n^5s`qY2$VH@AK}iT6yzkDI>AU?4f2>yk^osM)PWr~JeJ z_Kc1T{vD7@OV#PYWga3A-}!gHeWt%_Ge2CNM%2rKz4TF~ustjjr!$5Fo15FO9fax& zYvlXxdnA^Z6<)WxYbeTI93)p1vT&QsBy_w83CoIhHQKD}eW6jPFx-U0Uy>79&1ubq zcczp&+s5UVlAYZ~t6rpJo3Fq+gkl8q{9MgBRWO=%sGfkm`qoGy?HJ$=0g|ISjL*jz zeCTc@F@-CNk;PL~mP^?;tLysMDYrsjBF5I+(!0`=KqI?mi0L&jJX;EotN?dv>D2$^ zUjGe!#w1@Z1F!4Gu|Ems7o39749)GLiP*{)+r>zh%NQgg!+^f8#d}c%R&972uZiD7 zzSRCX@3&g`Pw+}@y+^q?rSh$*>6rZ$38uXr<2}%AQ#(=4^-IxS!DLt?VLM&!rkzx` z#LpRKVZM^Ray)Mfkg+juHh%%xaF&FjW`g7WFeoVf+WnQMRE85v|x%T+i2N~ zeXT|3z`Z}@0yM7x| z?SH+Qu9N8Jh|bN6_Z2zDN>HTko>2}XOSGeJJWi7Q@p$8Tg{I84(aBXcpmyp^mcI#NI$alLJG1TyTC-)lQl?D>(DAA%){cMEHxB7a4 zz)iZ>;%b|Hq}I%+o0zBB@+!iRcb;a~BqyXk_rikKL;UE}P<{JrHcUV}%_{o};I&2k zKIOt-w0=HWq%_J7Hfw*&;Zpqa@(JA%y{8{S?{_hf5YL{hcgFN(=Ft% z3$zu+byeRlCD+$1>NX#2g}xFKe`DC)wT<^a+T!U0B!7&`FlcULEO!LfN(=w8V0*J+ zt@nu$-Ih|EGJ#eC1@e^}lhWgSm({NQ3&2jV+wp0aqmBH;LD0oxvjS_3ShHnJc6M$d z^esN9C^s4OK8>{6`~(FPnH~jI9L3bM0TMnxH4-wkRjX|)8cV|4UC6)-XSNSa*GcC$ z>FIn339+qykJVtT5xJ+Lp5XX6xF{-;#m}&@Vd1ycxN`{B1F+ZAHNl{f&T`gno$IFu zGKWKEv6YpVuJxTxojn7a0D{-+{aaG}xFg4614$N(Cblnw57y{T#_AVWa8iqsU_*sX zwlp+71|ZacSQyr_LQtw%57e1|`?IKh4dypdWOCFeIMF1Mn|hiL_a)RkPn+`L?=ZfD#BlIY|YRl+S8-7akZB zf4~&SN}`a(r5;x3_QA2N%U4k3Z}7+lDu?M{@&;q=NIUx=H9)T5xc?V@0!54Gq5{#{ z%Ai`H2~FTSF!_}x?`3{#VpY}cVfvw_l50O=qL(E}K3)Cj5qW?$0|*)2C;7suhR$#ALOrlA*3q(7z=G}$09O|UAayjFTI*`vse;v(CX4Pz%QMlg^ufipjR3#ZTc}KYLX8a*C4BWb&qTa@URCuNPY7S7 zdb}VJ={Qik&&f7_bRO+WTJs^822@8Z?w};TEY~Z-Ig(VuY?gN=s}4dA za&wQM>10+vBDFN~g^dQpfQnWt3%Bcn?_RD_iU+&54clJ_-3)j&QN>T14y&k2#`3=g zB_1oqTM$2(ZBD%j8Y>2T8NWcH;UA3ktyiMl#1$$Yg)DYAqZ$uA{%`II(4&iw`^U8d z4WX^zcdLx6eoa499wiDXu87|1{x0sG_NvjNJav%*EYbq?N|JY=BQE~H`#qPst&=~2(Vztv#x7LJ@u`OV}hS9Ob0;U30p<{sx@{AFD z)F%HUX>c2QEgV;(cs~Im8r#gxsRXFRt zrKZQ47Z*tqNVxMXN2gMGo+~zT5p3~;oKr93S6oi) z(3|v~JNMhkzSS4KKR?PfU#U|Ojcv4TlOHLjbG1zxN;tE1cWRKB|D-ImRoCJljha@& zQY5Y^ZEc_cDRuER&k;kMIO`Zd>MGyqs}w$uzbZk5k!79xN_Ntz=L4py%f?ccc>PcKjpp}&(z=ps5= zycj%~xn_5mC*D~YT8$cb_x0)Mroc1h~FCUB7 zwnNmc?400VO$`?pz4b9IQ*OUq->6Qnf+j^=m&n;gWa#|JR?TFBkw+*{GAl!Iy2m>j zZu2cy+G{ukCSnH@Gcrz%R#Jc8Hp^&%X6qoGDV8Dj`kWG4T4UXD0oQpR7cU1q>rA6; z9O!V4mWF=YT^Z@OjP2Mdz5BlADIy{r3fwN9a(Y&b>(!sv1k~5+-o!QGcH7s{-4L-Uw_p-@A+e{5_m|Q zQL_r_5xbQMxhb68!t;{Q(+lm|h;lkfzKtStYv_WoCl{?;UR=HD*cz7{uUR0B_(+wj z!t(l;*Hd(QH@5`^IIPM7e~y+)9L zP+BhGM0?{-qK#;JhD}C#-u_f5%zJ;6CiIBYZ9TMAIcX#*A~mq!#|~O(E(6>41^$Tp zsQ00f9#yV_XI`JwmP}T21?yQL&QqGtb#R`tGV+6 zx3t+q1ekm0O@u3w8aaJiTl~4jeP-1hffW2SRt)%~SA@=I*KTgK0hO7$V;6dDRFp;R z+oKoz*P>#j;%xG=a2HbRNE5<{*V6(#T;uctBp_nEp9A5hBvWTM=E94seG^GN60E~; zhipA{0~&zU#z7PYW;GGuRzXTAD&J% zkdYad0MGK0eF99%82wQex~$|WIr z=*dNZ(I9IJ4>qW&_XJ)5l0Mm#X!5T8{>gVY3)*5F*(oD zPs_Y0sGhBvF)a@VMtde6SeMT$ja_dqKsd!AN$a13S6emXwZYWJ-dPIL$vtueL{q*i z*6MP6*BN@6n6>X@&r5~~uqg!^RI`j~KV4DC2;rbh0aauVoi@sV#TLoyv$-Xc&@<1T zK8zWBskH~YiX$v}Yk}bCNU6Xcmk`nRdI&_gL71DR-=VYt)7pYckRt7Vhf&xjg4QyY zWzJ&UfJp|*pHhGvP}64!oTw{@RrlE!D1q2WqRvfQ6Mm6~qVwt3!W6VlL*YS@t>9>9 zvQJ$b=b2gMwK>nND|78iQSGP|6smpA1mz%|C(PRapnU1oxSVT>$ri5!Z2+g=kXWvVt+9&DZ*C8Yg&Bvzd=s z{2O7%-Y@)7vZNf`x3E8uaDX1?W&g!w(But&$Hs-1{9a2P>FjT0 z47cdXQ@i?lu7XMix51#syGxXLBl1FtlpI|Cl6JhyG*X@)&CQCi#~gV{2>bwJjaNQY z>ba8t8f~xf!{GRZWIOu7+7PO^V*-s225I}FRUzY1I@JtTmXsJ8##R5`#gq2b$~aeD?TlBVjbEk)F0YiyW{LtC$bm^c{tEK5_t?*o(u z?ifq92y|Zp(f>S3WhC^0G)xWqBA4QDrpsW(zei)szB%99 zs|F6fnOM|JwOfd>f5rtH{ie#WupTotZkIL-x};OQsIl3PIf8Wtx$pav>~gQW3Fq3N zV?jJ?tDu>$6MHeNd_3dyj?Fys2?cE`pDd^jt?1p!%o=}8EjYlIR{M4MMk@#XYqkvh z&9?UX6f0Q51YJDS`jW)kNw2M0V>R*Ktq3<-bPcLf%K)aY<$H&P)FtMFhF^IfwF)KZ z0*IVdcA*w&M`OtNDwgYZN}jf1|W$a)T+3}F)>|~(dhgk zRAf1l6et~LxwiP$47(?A7n@a)Lt!m^=|Wn20WZf5%(a)(^?L2jH1)4Uf%MTn+tqh1 za9Cz(@rt5g_42%(~}&@lP(}Zyz8WVG`Eh(V<$1Htg4FLYEpTN zs3g;oAy9v11KLuiqcpg2QQHtd(9c{)L23ok(z+9jyP%@BcXEExQQc1uGksOL&Cbl{;wHnpaD}(NF=iyy6iz|C2Jz_y!qVhwrwdpi5;MoMN zq~G{pO{es$Mj5Oq8i8lh3agLha+P5t)X7u1tAmk|gI5hR{+V&fmC!bn1w7GwBX)KC$PtrG z25BmW4zux!4>q| z`P{|uHB!-h29c)fdy};{&BVW+J8+w-E}4|(Z3kO$KbUFsOb70oFvX6&S`~!+>LZtG zD5bD=m{vf_mHD`P@$>r$&i)G(Cf4_ucU$1(aPX|(#u&;2m&THOd*jGtCzf~^qy7Ei zgEZycktBgSwa5W9*y#$W8+i0QwJg<71E}(sFTO8V#&XAY z=@7fnu@x3o)iSGV&C#E$mw7tS2bdS1(@g8L&c9c+TG4mH}h(Ly&3tPDEnS3TLVv>)0@Jq z&2+f zc}5w7BfpCdc}!6j51rCg=;5?&H*+Uwk~Yfm@WQ^ZfJW#Q>I9y%_>zI{}`t8Yh+bt;PQiF>I!~8y?1*#iW(Z;_v=CXcO*Mr z0V~&7C@ifbeSKWcp3H<-!|s-O@`f>1vos%ub}ddx}K)rngQbn2%r>e`|`e9|S&Uf!P3(%a_4%TlN9rWJ<3DZXg?2R)+J*xXS z@*6h3O+|bD%L@ZTGXlcQ8yBw}@_NVFw;*C#N7qB9$EGB~KX}8j&z13dx0k{TmlOwA zdCugH#8*jvnvhBL>({kKf-`QPzORL-9c-1gF-WVe zaX-oEy}y4LG*W`;189R+A^{^>nk<1jrwJUT&%l^$bx3+7L-dEKa}IRsmNiRge{yU{ zOkf@4wyuGedj&PbL&Q?@q}+^Nh>BUq?{PGzMsV!??&&V(wFYgCnLl1I5yU2=zi=5$ zU%ww)&A24%w&3c-87*5A`wzbcRJ}g`s(LFSc|`}dZwwT693A6x<8Mnzv$PemzJC(w zoaDv5@STe)f8d{>Lpslq(ZNr=UJmGgJ2h)!9qp@Z{m|9Ka|tX;9b-|#Q?l#EEklF+ zR6&nsoWyV}Br8*^^X_-g!W=5!0-=%*@=w&4uT zU0N2?QYZw?mB?Kvz>F{r&hMr>lb@@C19Y%Sx7!AMmcxv%!b`*@OWtY3%0nHhior} z5`L<|Faju%v$dkL>eB~4Mj0Hp+!{lx@BhN=UH01yM2!58;aeT@zYb0M>MSTrx$Ijc zgoX|T%C42d2QPlR@q|6IVl7H1{B`fwH=F$632UJ+;Qz`fP8~9DN=P@kcMNwaOE~k1 zOFU2Rx@%x8dQIiGxsi;KC072c{k3_9oSBq+Q}OLzbT}Jv!Rq8wcS%3Q2=XIc;&;+= zXTIi;mbWO9!f|roFQD*K+OyLYAHuL;ph&ivr32EQ>zjBTWp|pD)g=EhJ;Y)`-3(KGsCbNPgDM@GNdcadm_c~tTKDebN}h z*5#~k=0J4pVAL*Wp(Kzj6Wf-W88F8Q z*ZreLoU=Lf+*{t-uE^V(d7qxMUS=Tw_bz{W_4f*E&hoAFm!!7C3R49gmMVYO0c#vl z&@shxnxTmd~flWmQ^F_H=&zgQinVdYr6?Y<` z29J0^`Z24(L$BE@sSlvGoTjD|VYxA`IT}gJ$0C*Sf>eShn7OiU63#g{fXQ;?z&q8M z)E~l_Rwfo4mLV4nms>RA>XliBmr$qh_9Yd1Yolp}B(fc?vFq_m2j-$8kmninGdIyQ z{bD(Xo$w0p@Mr}#gO&>ai~lb9Unk<6fSqP@Q7NKGk7%csPbvY43B;vCLgK-o2>BBr z;+L76GaBERbC$D4z3O54Y{InNGT5>l7WUWm>55K>$!g_KB3jyFWteZ(LX=oquq)vE z+Zi+o}sh#-n*&sV|&4VD=s8%Zx(;%Gt`|KKDAlQ!9r-6XU0C z6kF8+M|y|yg^+N6dQ27uB{cgMrb4P+w-27S^DMZ=oP@Tx)FpUo?fDFsLI=}xhc^`? z6=-v~@-e4rYX-UWF+jB^a1}oi+8q^>n|6b=CB^yoW#L+8uYu z=+xGrsZ!hdY*J8<*$*F0TvfuT;EZeDUM0YW3><6R^m$CIAs|;xlY~{_<*;YniKICE zzD)E5f(nW9o=?%XbbiMDQBPfuMRuN7Mo+KBi)a1G{qd>3zL6eNJ~#jtnT>IcOwyGm zSOQb6#loGNnptw+eBi@V-%=-TB!NY9FTh^JMm1-fewNZ_}){ z#bXYS_@_yi~tW2iwXuEbN8bik*=c#aNq%+RdXYHm6UwyA!af+^PnvbkC%hx_n z0WMs_q^gHoiIQku&rQCznczurCbABxfh9yId4@KI=2-<4R=P9g>gDKtd7U!A>?F z^Mg;VH^N2gbO@#4b$rUrFsRpv+}CmjUe6whJq%`X8nW0(rK{7?-XlK%;56npv97D!^Xwi}Dz*JPk z*!AG~)ftIK$zJH@{o3;C_&5hGt<*o>35;k(S%GsPo2~9iPn`lVn&{R9UfX5#q!IT+ zo10D3Yc7;+Az8QhZE%wmfW1Ya!)fD(s?7I#Z0a`HchZD~ha>xF0U-7Ro;=UoBk8W9 z((m)YsKoytdjK*3&tx5DfOeveu|$QOKGx^&wEPp+J%Nv3Ih}1r^tfDrbpWgxb}l z;G5W+z<-B0BwGL2m~Qd!4KcCKpSHNXyva8p`M~l96tGJX|GwXBlcpmcNzul?MGad$ zv%#hRqaIR6H`P7xu9RTF^w`0wsz*L@u8NNDb5`a8dU2(4#`&t%_weJ;e4#->pQ(JFK{7&x?{x+VZWg zj5*!S4@VR=^!D~HEG!)Iz54G(7$^G|7ipu$lrecBhOe&30`tf^DL3{nY>ezI!)J56XG!GlIIS5ANQZShvAsHA* zD-S<4p?_#3WD%4-3u{4-74V4f?*AzvcpT6lb6AZ0`vg6H(!c{6pjs8JFo>Z`HTk$LA#RoqmE{6YQZR^fWByRTYYO}uP~ znq)QON4yWo*M$ap)w(Gm4(~DtO!RcaK;pP302j`B ztFbK)l6cuyb_Qp#_M?o&^wiY1f1cuh|JK?6ewfR(o}XLUJmo~^=BYkkb!o$igy7&; z$Y?i9BJ1E4y#W5vxj|urw)e*Gh|`rz7M3W?tT55LJaK^Tp}Zvr``na|@6`g7Q5oYs zX9PqF((w$TYfRYYdiv2@!%#DJ?Aj|{=8Al?SRp3B?QT=Jp?cvT98 zEaBr8=g|4^`-}8J(=yb&E2)v=%CCgs&&M1bK9tU?vxC*hFQS=csBD3pC--w`9w9hr z75Y5|am+`c-r^yn5QTB354Try^E7Dpgja!D_>YW+tp4B_<2T#Oo>5~E}t|yZ{ zU^G*32_U_3=31#i~}$6PjKz^fDds?o1pH>p4}2Ps3C8Aceiv+;6K~(e}yI+ zbFGT#ln_>zQ#G_g`ECYCO`G}k0EFd}=pWnSu4#Pv0HL-#A3hNm`HL0HF%m?|OLfSx zp9ZJLiUlRPUcIO1i6!0E-Ktu`8)47Syn*ZL%LqVRr-cXb&nU9Mf3_;OE4k%AA2{Iq+l?KMiRd7 z#a%kBOXxR-h}?oBmA;)wX%0e^80;r)I0u|=2gd%zpR6c*dtVU5crgx{tnNwwu5>r& zTA&fn7Oblir_5K{T51jBgJQcCz8(?%R>j?Aq3B{#gN06CJ{;{XGpr!Htak2fsbf(> zMO_Cse=my|?VSErqm?@tD|xQQ%qWNUW_c}@qF*$|qN;XZv?dx(*z=@P@1 z=JRQ#+=;QQ_fLS?THOuii^Wr5{R1!W{!GA@Joyb+lxfq*vSjk^8`#VLx)f5DAJceY zZ?V;MJ%Bv&;%?_S!0=?V?Q`L!SiVKo>&U2E_<_2L$~XVL|8CNM(%JG%sve1)kNTPl zmm;pL8($GYmobIV)QB2a{m{kMQPer#Wndyks1;+8yJf1bP_U8QgzXQb(NknR90G^U z`2F5w>O=oH6Kfsm`2tQDA$l{}0cg_cW>XbMMko$xP}-3#G!<0cNJR}l0AiNOGZ3*y zT_eL#L0N1)Z2}s7iVKuRcHMZXm>=!hKItG}FeIf7Loz`8lt$1dOHE(^qC;HrK=6`$-czo#T={fYk{O@M{`)wH-WL7O-eiMOHFhALG z_IqWNsp!-jfcQc4u*%yVQxUCS zD6_Md`RQoRYXPf!GB51j#O4vkY@tL{V`n0Ou_#|YxBGUW+{icIdP;q0Iy;QF1<563 z*2<*tCil&P<;e9fI5}tcoX4Lfq$yUJ9%$}J0a$p7q}o{@MDrUV>_Z?;CbFP#a|0xv zZ8^l7m2lLA`URDL>u_zQhWGukHC>SSdiSxFk{cWaR?}MWV`H=drJk9P)AswyP5w~i zRM6Uvbn5`-m^^Bem%m-`W76~5%wpx!f2ja|b{6BLz}*Whl?r`+#Vy zdm^JrEWkM=^WObu&Q!B(#D#rtfu`8GCF_CZ+%DuLt~9vPbsYgg@e5vv^vjfX)U@N> zwhDlEz60a70YpR@f|yCGY_q_ADILSkv1>n^3%sAddKvx3 zbkaY^DrpFyR6z?B>|ab7dZkn`!*M_#5G|b+lYV`gI7w~rl-k9TELD0 z=x$!!x)n~=aMS>E@duSE9`+BZuj~A}_@BhCJ8SD>E-h=U*IFpdCX$W>J>%;y1mGYj zzaAm+rs|VKV%(a$g3qd5cAQa*dBztXp#UB-RQzFaDGY_gvYJB_d^FuR?-rI7S*uAd zBl#>ouEXr$p!!C5bNkE9$WjHMby!<#rZ-9HOL`)0QW+PkGC`oyQO>F`=M7zLZ>_eo zYcXVldHG``4*YP+9FTBx(^0^ejw`50-iPN(u_hKwUUW*!1{R-t62{=zf_gvVNQY1n zqg7WF&M)fQ+A?%HT9X4{bFFu^%YE!PcADVNS{X2x>$te7C7Ik~!_8+%cQMFQ>Y*Ia ztpMXB*gUrTlRP52MsV(qqvM*^iCgy zJU3yw$WJfJz;L#j-8w-obg9hJS@~IgrvOWJwYm7Mis{QnF79}!R|2&SH%7{cKKJsq zq|aDq-8#s!ZQRxN*3%d~7*IUm`gvqfI$l+$SHrSvVBd|vtEBl-^5H%Dp~KD@H|XLQ zW<9KeIsPvFyX`iI3|XvM?STj-%4@CXnkP|haeRJiIgA#g*^8&CYy_$aaOD z3_ByG{|;3kTYgU+8JP?Upp{(6>wJ5z%Bbw0-u|+JI@hBsQvR3>L3nIadfClKm}z)x z3b)y$EoDxvTNl!P8tTpJq?*D>FR1<$dSTLR@$A9b%p>A9CY9y)lSs&jYN^CyrK+E8 z`0i19RbATy^syjbx~K2NKT3^~c?jkHetQrvQolFmc%iJgeh4cRN)}`yVY!SLX<=Y( zYS?zbc&$`n^owizc%x@Re2V?CCXYg3-^DQCzVF_ETvp`P_M|Uw={$gV7Yn}f%&_WZ#L!d4qEWFz4mH&^= z>Ez@j=~81Nzx~O$*r9dyZqOwrji@4#CupH&%@K!sBfDSc&|sD#?op7nXF4a#@%KdY&!pHaa@ku1yo%eVTkA6CWd_?+QUkmBvcYSC`*?JN5mM2rnq ze|zGt$)U0sCsg9utDJqw137?`xnan!+)LG6Li6&}_$Kr(s0BbzT9ZDaaZBQ(9cXhV zTW?ffzXKw-gyL`=KpufAba&5iDs0tU?VA6NwF{v*Vt>YrX)e{^&-b4K!?Fbep#iZB2}VBUgf^cC>- zBl|}60@l5fchxU3&CU{E4N1I0gI|(nhCI*q_hzi!HaOq}q#;=S0m%Ta29(@xuGpBN zdt81C9hi9PUL>71598^t!*>w?M#SV$(-q3EW5PeW;RvIOu>q}D$@=Q%JKSlKvNor_ z!No=T!`A%!_qqLFf9==kMnI_Yog(>2-&aofmacv|_R)nhAqx6`jCWY2J|0-o8^Oj{ zT30n*kp8AvlS)F>0p32z%wbM+YJod!?pDDM3OfninFm3DY>Jof3BXZUY4lH?In;At zVZX|zZ5KD z+y?M*)jK=B!FNO~On@&NYq+hqk$K7Qf%W$X?z0D5MJ>-%pHt@cOX(23rKqoi?ffzFS@*-^g#8_1qWEe1ccfFpfCN|q#sBzd~(~j%02}o4r%kZJG~Yk5E2ik zD=o=#h5vl8QBgmXYg6*@WVJF;u+RtCjG*Pi zGT#K0@4p}zO95WCyZbbtT0>G7x()NG34;1-7d2wM7M>hnAmpP3Kng)+zvSy5rzte2 zGGP&oUPX!|RZiRX=EQO#q$>+e$%98W&pmzeE&kO1S=)B@1?oofrFN!BmVgZ9B|DGa z*__375UR^k_ER!G8XBbCe@GIT1q2ob2ErXa@=Vz1>C*#c27>|7onHFNw|L<%4+XTu zuMNkOOS!9dJ9qn=U*vHMM@u?IjINUbv`Hehpej2@VFD`gYC1SzP z7~5k@7ip@YP;-$4LTaY;*W-g|Mb}1+%zXf;t3(TEVKyA zGSJkMop`8gY>0R1WR706-hW_}zro%kfhH?6Ga@o_#)0g=r|;iEY)(U6-30-+0!|VS zVavy^FMFEl$}>VKK?80_*9P5yEwaEfeR^qwv1ny2GA~sR`dg_ilST{LzW^uue3nSB zEK;(1mlt1Zze;S(d#y@5GsQW|eqx#J@zHRgDG5>0hvOwdO+rug76t#Wzs!YY--&CK zG@(~}C6ZP`RZ+H$r;H^rl266(TH`FK~$o+Cg+5nK}U&Qv}+VuJS?`Yi+c9zIhj5_$3Vwa!tXOL3AA zpnAP**#(GHUOtR3cW4l{h&E^MWN0`xF(rVo#={TgGA;7&Ll{iYz?ZIMk*0*=SRo-H zR#sO3h?IH*wAd=l2be+WDDX^1IY#zwcYwM*Za?EWwTRX?hQO_+FX20_q(0YTWpeKw zepA(Ot|a z$zDj5nmV;o`BH)fzqp{FvDN8M3_9K^cZw)G)@#& z#b4_0w6BM2Aiwx+rwGL31{{yeXaz@-E~mh*4Nq3WD-gkp+pl*%z@vmql$)?_Wr7`U z-1J5CdQMV)tv~hwQ3{xJ`uhPHN9(oW_{wVRFn7BHT`*$WZ1RBftw2wUvE;ONdNxPVEVFk3$iod7Ta zE5*K&;*=xod2{|s1`-BgODR}hsQI}2S0;rjF7-=c+rI5m6D#+P z(OGR$JG;IwY$0?-z}#LwhgFSsum4`&&sTrzQ>T_19X!rf6|YI8uq4%wE^pb~O*k*( zN?l5|esczHzYKoS;1{}2OjxahF6ojX_2u@O-lAa7uq^p(wjAdz1!;_$I9S!lGQ>y_ zx;zClYP45KSO*B591*{7(S?+f1ON)#BIk}!@y@q!xt|5fOQgpX6N{x##P5=3BOH?W z1Eqnhgw|l`Xwsm!g0uhw)bK7n2bbaHIR^;0S{p?KpIwP9O=uh7+F&Mgm*uVP+`a6{ z&&J9^6|={VAEsQ-u$~=dQ^oH;+~5A=btx?^74*O3Un#?sLXRG9rHY|#Rk0S;JM)}O znwUfDW!DZhO$}G9vMPJ#BM{r^lt2S@zYn?T?lx-A_?b%qDDr2@iGTs)xP=lQ)L9r6 zFc7@P1Jl2hBt*mW7-pu+S$krx=#WM`1Oi1X)f-6LiNU&6t13e2e#&AjsPlbWFVii- z1&h%B{`n=bb*ZZ!xf}8QJk*139wxKo8&%*D2J|Wdm%LK*G#S}(>Fds6FHI521<(0X zdiv#+$o(h2EvU3pTFDJmNVLr-L5dqK06ZNE%gr^`)lL7$kal=<#BR`B@FpJ@=P;6v z?OY)4j;{}3H+`*5*BAC;f@DkyPx(L2%MA=197(bq-*UUTy1@xem?NXh`j&;;uiPxO z6~lq%E}%<^k+&2DtA8&~qyKQHW-xDG>s4c4@R&2v;SUo|9hRB?N%;IEQiKUgpw9et zb~7FSRd$}sts%LbdhfDk6YD~^MJ=R4Rj->Psc&a=?Ty#z!Bl4_H}Cm?J={K}Kur!fG` zxE{RASW$yG9J9(en4(x)n4bh7pl#vx;o-9G?hXI6BsH&baBh4>#j!}(s@D1$Yr0>L z^c@9oBPSj;A!ZUBS_*I(f6wdoP zJ^ji5Gvi`nVj_9O_2JLEU;f-W6SlivN?Tqy#H9k+=CwYfC^17=iT3Yl=Buw%QY&YD z`Xr%buNgOH=(OoZEJu$@Evii3DoBg|7zw%sEGU zJz+mnsYkK55iURk$BNEiybDBMeUC2y+ZjC*iQRypx)2Ay1Bvqdg zBK$1e3l4AR#w|~;+*`I=c@7&vo*x(^ikDtPrDG2uYEV!fyTu!&NZ&J^T&Zg{4YoYc z7ciwPj?=tj5+Oq2R-?bN{i1JNH>R{gh_(JXYdsgM1j~XDGPJC+EQcGcmfFfJK6AdP zytt5P0A`pbcC&Y$ykt$(7>D5_!j8T4<(2?}8=z537Ao2~slV-X#3@v5=Vg7;{eeB| z<|U+L+H9UKV$L>$prYgJZpL{Sqs=6Z z_>Dh^Dk258j?--Dre#IdYedF-+@dO=`yV%_1o$~j?>on}0l(pBY%$M-UygdY1aww9 z71|=atQ_~~B|$IX7bPhtfTk_XH%@b-jkpKRsVh2UE!>aFQRYYhbd} zG3z0-2qY9EdMG2ut{+2vbTJZRlgLa;Z*E{~__giT!NFKhPaRKa{a;HqF3mMbfuyHz(RrXuRxH?Z_4PV9h`g0b8JZ zD%Gi`a~+_6=+g5umE6xMv>QlfAIieP+z*v;Z^+vY7Gc8EH5R!a zlrIZvmqk-0w>2)F6cntUpf@+0ylXHsTFc+|^Of!idXP9ieSN=wN`|6s+?ujui8Nja zUVokP&9QHuaZz8zBPPW;@A>gNnTC4O+g|`}2~e@ibxCyT!}s&2E>n8pJN-p*7@V`k zG&*Evw6{Pax#$FZk*@9#GM`y$Bqf&nilAp~UfgTYQV>*+VzJR#tA?HWQs|M3tyC+e zh#)o>9Wt0hskxlY@qhJ-6WlMxi+-P&i43mtR3W$WxXDXYpGVTyOvzg`v_!qJ8=99o zEo6F)V!vAR_RGKPb$0K3jpY@^k&5Kn#jA~M`G{06!!o6c?E;4C*H&eRs_W^hqvZY^eTRdvcz`4?nOD=eW6*CdU{UbAL}AzTDO{bm@8muD(8q z(o?7|^&XJDt{{-Vd?rkH^Fd#ujTWr>{~sF#2{O!lg-eNltaEvLDaSd9dmJt%qVNaf*)_(K_) zY=O6#lBm)5{FbKzXlezM{e>< z;)I8)ra5<~$mbpt#c$(98_W*lO!fm*U>l>84HGl1jWZB)vld=KW~I>A#S&WW3u-nW zdSs~;-8)}@`Kq^=n7HRYLED`mu&UY(Twg9gS%3Fb^_hp3`}4XVy{Dl(`5hK~jow_l zd0}6mE&G?p<0Ow-$sc_;;UWw!idUp-a{FDHx9vtRxv7#)twARlo?0QcFL_0vG8LVS z9GmhrGXB<|Ga;Zd-0nsv(XmA4BqS7jWT zm`9t9^6;_L7+n3=@9M-&vmkpL78CMV;au+K!o-e~mqk(CyDyx3wvKJJzM zThD!3-j^K4S?k_Wh-k>Xidg05%4X;td73c#;c0%SoE5u&=BK$^L=D>3B8rsR86iij z9}m7um8}+)mBm3Ip%MLq0|ORHq-X~5veoP36V{=|>6(n`Y6EatnLXV~d|aTC5((7E z#$tY`&GP+h4A6ub=7>!e>n*g&4lPYxADQ%&k_6xzl-I`4TP5PLqZ!GynpglQbhX-O z;RMsSj9oDGFfUj!TJEOX7Phgk5vQJwn5d9Eb}5*Crv=TCh#4wAz}r~Zm>3#bngicY zh28F9NrMXdN;qz*kX=f$Z$;$9iVTM`KnfTp0y%CPW!#e0XIt+~1Jl0W;68Z7#!So2 z#aFZ?UX;)GrysmQ78L;oN<{dvji@*955JDuNU4$8)+uSwGnCE3u|U90^*O7uw# zQ5{mZe)r_1^_Y~nX$>yk5Y6rTTf|O56)ZGa$TbstY9YY3ca8K3V$?39=7s(CQ(r9XW%srfJlz&e= z-Cg9;MsSR_Pj!wpPUM^6W_;vGABp!28+8!knLG?fnWsr&%wz>>XQp@DwWO($V|S`q z*SC+e-hIjQc|5iqW8`z^M(?#}&+b^fuE(f8zS zADLY)LwmMD6}wwx4R1={lt}G*{PtGU+oZ2v1AlRYtU-6>Z@=|BBMAS~DfZj)jO(N6 zw<>5kU6c#`-l|Drzo$)NhQR&&e6z3FF*80}pXN7lgnMfY{`!Q@_6ZvnHB#8=lcD<+ zv?CN6$!SMse%bDChD*uPYF-bR#fd0dx>YkW*uMW$K%j!(*}h~$5V+@YZGo}?;I@`2 zyqW>^`d7yn541{*@JRaC`N0Z#O0Q9pN9kt=D0nPS7l&EQRpblFDQ-Vo?vhgw)H>C z;N^DBn@W)hyHgK^dqe|8Qm>V>JqX;rU@8J~YNF#VsdbIv1(P1dOW)@Utwm3Z4tL;J zF(dAA zo9#T_a((SEd2&|-AebDu#NXtb^j#Tv*QQH&-tu>KcNgH}%l&`UC$x@^wkE|rj&@EF zpZ;_eYxq&jP%_B`Sm|{`t-1WyJvAhiQzZYK-=6)xqz2~uV5Md6;>3mFz6W%+knmE^ zU7Um<%8FDBHF3OnIn#y*x>!rk4P-ViVcLdonuohT6TP^~VX(dy?qB&Pj2~O*%81|$ z6B6NkkIT{DEITw%AHAibNByJRcz;YLv9)`a%kqe3WjE!fT+K1vrs1U7=PGR7haaEs z_P)9YBEB_hWaQ4~exiLRs>iV8?VUqX%GUSu;CG5-Mqh+Nq<=gMu=bILfgcMWuq8fw z^L<~W-2a4^(ObQ4qpylVa^R5<7-^S&>HxPU>(J>{I@93GOVOZTU8HERD?n8=KnUHTYR-1-F zK@>G=`~sDD23h+t-skSt1Su*hO9KX`k4%rQwy?{I;<+|96e9v;R!w>SWVAs;nxLpo z3KvjTK&cV7I&8gcb7HaWV`mG4|?sX zT}e{wUezOzayP0Qca~*rLC(H;NNVoYbUi?*j}bhw?4vN*W2;loR#|K7y{Msr_NVq< zym3XII=)k#Tc*E?6g;G8D@F0!K8=kzcbH1I*bO*T?&aX%2nq_Czj^1E_hK^u8soFH zZEd%!99KkKHgo&UiCJ>0-H#HhOb=T~-FgST71AAgH^zeuCL*aEBsj$mO-`z)qzueb ziTqMLDczSXfB_R7XZ1|FIiIx-s*|fI1q<^?SYB6RJ>>lGkEqK%+e^ag%nG-K*~_{D z$La4V+K$2z87{aHyHlg&7|A=>RhqGME9`PIXxMHGT2p1$lP!U?Iz1WX0>5aVtbn;O zrtVQ7#A{hAV~%00_@(WorMEMK!_&hp8k~#k>p#{DVBg>PCON5@iRP6m=a0Ve-5e&% zYySM|%QY=+2y^7v!#9>FHF>%+sdB%yEn|j4OBF#sRX4cCZM$>A85gN~iQamNZLfL; z4OR~VO_XYda!YGiewnTIWk5N0vdrc{`W$~oVlHZVq0&Iq+6&=^za3gXr6pTwv4NIv zy&G!?u~MFadTzgo+qsop|1d>yPTP||)_s91C{V&>E|1e6^Sb=(XjWddEgsq{Xnh<4 z25o&$YXS)E76Rg5yYQ}&i{dx=s?E*K+}!6XDk{jiK}~*a;#JYXch7ynsl2gZ$w6D^ zD*_rKznuEJlE*sFH#bW^rT8bi(YT+bQm~Xrh35p&d6X@FPi7*pBVw1Z_n51v-G}QZ zz1+6LC8@S+i(Jpg$KS}AS&u#CV|7{CjO!tnT3{{+z4X-R9@5yz@ssT% z42h#Hhb?y^jVRd|E9}QJW%hG8QcMQAe=rR$v(ro3tqB?J`$@(sCDj=o4FtMn`xWRG z2{;TVWb#Y(L(|JQLuO&`_p@(81sNwO;5+Z9(w@?Wr|5aJIbR0hjFfdlgyF2rhqg}Z zO}Rs%?L$e}yz7Fx9kIQ}u6BhgE6ydf=VTT9)q${RYK6Tw1bGMjj4W|^`cB0C$$T}- zsp;u*+v$YZSQZJ$ueJL|=AZYQ>X|Ois}7JIy12N=`!%*qpU38*8hgd5a4X9CkmXdW zvyZ`*?-t3Onv~A)C&s!@K-z!>JExW@f{GpX^Z3lK+-E^IM9@GTa*k84_76suUoTfB z{`=*1uDkOmy9@LpPT$S>K*)Tw4H`4RRg`}ddHyPM;sJ_Mw|azHXP zH=K@UFNNN>?1tX$RbhGpvRU$@Q?MI_*Ony#6?8QeKMk_PayVbqWz+R^FY9FN*N{F4 z8Un}PqMbg(CZ}DuO^&I$7X+ zFbJ&)#pu=bBntAr`K3s6CH*-&Yq#C`mec(e0uq|V#l=k3=ZO4~6O+zhnRDBKRF4M@ zuaSYPIAH^7QsJ1Z!CY-MlwGvE9|Y~<%x_Vg(XX>5d56=*3g!?IAW&I3Z*Qm?I|&OLE`@8$yrHP$ZJnw@8hr$ak&>+z&E$@ z2G{bzc|2ylNdcsc!KDAZ_Br>@4W(RKTAHr+6crL08y(d`YqBoF!XVHF+}W|Gf8Q9w zYpG+isI>2inxob}twvB_{f{TTE5OMk7Cm?BXz<8rf@|hYzhaKsD1&HejRN5mKa0C= z389+*j{s-p_RqoBl$4f|laZ~DV(jefFq$Y}BNUl6EVh29#cuN$wl2m!RUnmYPtC6n zz z%C-NOUh`Xx(a&9+0Cq8gQBvalx8E2U8PS?nkMckdVL`7R$4;yzd*%zS3mKmkeeP*U zRAv`P*bR5I&8?)!CG{;;pkj*mBupNh`lMgi?(B0)T0dx2Q}AAvJg+gmg4Gt+4*LV$ zP2GqZc8xlbn}^+l6m7gzTPleuuGM{^!#!R!KYr-_l~WwmWkQoRFgaP{xYFAJ6g1Y> z=CS|Or|}^Mpx>gM-Q3*br&fdtD$DSvPOoL00G{W_8+suq4NbEHaT};1G7T>u9LKPo&?3&*WUf zm$#}ZYu7ft2M#g^qtZ#kpfxMf6&_9}M`>MO>tr;pzW5@fh?uY5um^D`492899 z_O@UX;5}_@^sH!se$oyh4Xogf+GZEM_>S1%r`(fU4vrE|_iWfPSkAfTo=t1)OTdRx z67pWWfRM1ESunh8+Vfcj?dw{`4SYJ1-oj+{{hd(}v2@jvR+CGS22gLBiXy#9+Zkkj zS3#Hd<@qcRl+~bWWX(xgSs5Uh6qS{ir9FRWik6AS)NsEEfbG?l~)}3;yTLOL_X*hM=8>4XHkI6FD$GC zoKsYc91rx~fh&O`<8d4`F z&)Ccn!Yc5KSR|jHP9FF!=0r3gqN9?<>tzCG>dv1bD({d(@_sGg#ygx?0To*#_EJg~ zY8ArS&=TCMd5L#dK5 z3U&&A68nxwGFc$C=spoqNfDcehew6Yqzj<@VNprx?(L<-Ec}W*C-y)L`;D!ur^oYP zwZFlKKw4TFSc};@k7LuG1PzSdbYEXzPfrivNrWERgZy_=wN zLl)ha=Q=F3&Yz!wfB*gW=Dvl=yCFHTv9YD4r6&~ZDLFYgI2>+g zvCM9^9$I1BVGV}2VAdfNNsKHN?)p(scCjrmTo;ln4kZ! zIfxtxJTrh`=;KX-aC*@rxwnl>1f&d8&{_cc0nx0j{;iYKK_`+`THbVRAgiIFVNl+B z9r(ltiweVD8VQ|H5%8;$S6XZR838**Es`B=+=!E%7wXx8`k}{9YNtY|I4SK2@KR&d zBZa!GGn?|v930AHA6gO*YyNPd>(|>?4ri3~+V&CgzPix&zC1iOa|Y%O#c#KU(awvW zfjM-mpLsDT{OFG7ZR}$4k_myY551Thevpy`jPh>C4iHXSi1>bOmjusGXY-%u{S6?z z{m)bM0kGu%dER{kpzQxVY(JOhe;!J%|D{a(jp(NZ{O9>_8-Nf0`Ok0Pe*gdL&qY6W zW0bNDszY#c?##+wkkxbU-Z|KA2EHovY^G+%ucJ%9fKZTqH}=jdx(`CQ}d~vh^(&>-G*4o>i>>}3a#VZ!Ixl<_dsoqx|#f!TOZ9rX(ltHW# zS_=f!E8SXFt-`zBdbO^*I`@)<9hbkSNsG?v%FD|uDt1?sJ_dmR7h?!0zW2LQ+)iWI zw4aFGn_pRK3JxaI1MrngkbXOyWvV0cDbUoZ-a@M$q+4nKSM6}jha!X7YG;L7-r9cP zWcup9sg0zdJY4Yre`NPePju#W?HK|zFSG9%L}~V5j`p5PU{R)io%`XB2A@_X zg)sph9vcJ{sl zDIX+dz{;8N<{_wR+;Oi#J-?ZA-VEq+oZrIU+T7aG)zxhgZJKm-a|?nO3GN;L473mN z#VvmvZQmiG351^;-c@+D(whu4sXf(`U|So-%mzMq+D!DB3E0ULQF!6no<-l6P}7xGrOVUMn4W~I3te=K zb9YMXFTVCP1vuVS=4!=Sv~ z>#oI8ml(B^}E#Q#_?uG??&}7FuT`fT#)wcq*sQ}${2{^Q!Oiv39T+Whbj@_ zD&ea0T1YWTNO?q#&&*qkjOdO<|0v|Xih}2%(-L_Mw%lftQ0~TV@jtG(<)`~LQ2;zr z0>2dyM~Qjw_^D;RtvYw(yU6P(`rgm$(yi|$g6J70Y1PyZyLWCiT{1YQ5FS7Sb$w>M z9PW3U{qBr_-Ci~|X5x_;x5JsQuLUMJRlt4P*K42Ps@(pwD$dWxy>ENgS+U~2boIlL zLhrUavx0?sDEBZ3zi<9%{sYyT1Qi>}w=?{n;wU#aO;-KSv2q_`6LA3wOo9#+(@BV= zLyDvWiX0qs2HKbj*}PTZ-``7=m+;l^#Pcm zTWYE#C@2UdeMvclD9G#A?Om~)Edjrr?r(q7OP!wqgo2-}zyR*{)GRG;`a;Yr;h?L8 z_rX1FD7j-flc^o~)f0}I*Nx6IX-Q9*maM%;eMN)%UI$fGSX@qr`N_<_39fxmv20n8 zJFI0Sbsb7r_sn*FTqfQWKCjR>JNAt5iXUFdc_HrgsJ}`aSs89reH7O}dC`HuyZ3I7!h2ut zYn5mA(DS*TwmWd+676K{Y+wVSB)?*6Fz6o~wcIHcd*yS1(2fxcI4(|^v zGbHTHo2A(`_Qlham2(ZI1{gtoLJ5~33>gMO$)mq2VVX`c1~UyQw{Oba+!AGA6uLaP zK)1m;^y(j_{s90`xTvV8!PclbUl^@mz1L;+$=!cWgThbms~wq`=mCT;AkaMW>PVH< z$WTImc0{${8)JLA1ITa3_X#wcP?OEhwdFn|_3tK=3`-Slx?^hfrfJkP+0qU==>yXV8?+$)f^;xwdU_hZJ@rE3A0+whZ@z~d27>`;7NwR4(!`YgXy5a- z-WA<^^$NTmQdN<+5LieUq=Oj^@be2I{Rd@oH34|D)3>|38^8j2L;0Yv{K9%NqpSNr zM|w8Tqt^OLrR1u$v!THNfIk2q_mGkO*CbUR{v>J0w`XyZMy3zAOOCgp+^8 zh^aha5A(JK#df+Mt~1fmG|LhF8uG)Izb892H3guteC)bB2`E4>T)v$WSEqc?5nLa| z=W>~Or4SGraq4**AK0F_TO1+KW|sNr4k{i>jUn?KnjeVrzN*W4^<-_V*hp5^|HS;4 zrF=;L+qfMLR^Qg?;k7d6=Os0H6tnuqO4R4Vd&31r1VCAx9bMy`?HI2~|75vKJ~o8n zPPki_il1B-leO#nN0@}!n82$h;0FlU9)U_@u;%^qKwdX5{a5)6zQ;l>| zsE+&Ix7*T7BPC~t8?#l8a){htOY`#W-=;3)vKa?}$C}u-?^#c&vj@zqFYCg{4Rgx3 z0p2+AtWd%h7O#WW6t0{Q6&kbRqY3&Z`vlnzrorkb%V%h+1Oi6x#q#RFt0jx&SeNVj zNNYI+YYC|T|1}Lihq(34r2+WdOaHtBH-+vQ=8LK;TeWK@X&Y;se14(?xV<*8IzG~} zWwaB;mOj5B^yt)LtEJ-r7QMsS#yEhX}3{bGoORG@FXqEgNtKZj~7#6=lA&+iJ3UGU(|@DE#;s z{QRK^OO5TId#m^aS!`tO&X_!|`U<$Z?}+0*j-7!nHpSYj?-f|bVGSx0cG@qbczKS; zwFLZ=6aKe3gF$leOBLKTO>DzOkc77trmcK?Qq#}|T4jo)$g}Fk{Bdr-zv%m3jz!L3 zki?qop^LD{QGd)2`h7a_!$v@@KG4+1>v~@6tGG)Nlv_DJP~pU2^(ec#vQh^~W)Nln z+fEb#Lce7M4b=c4eL{?KiX~gb4FO}s0aw024=mYPS$0YDYcCWlPbqw#T;960y513_=S(1r?53fiA@FYh|6Wd>Ilv8Y zb94I};`tfS)B=ligW$0q%f$<#5Pm5Pa zZ%2j50r+Ik$bYiyTJoR1oY~mK1g%pwu)~eB4cBFHUKv#!i2XjXQOxd9(ZQH~#wkn< zSp=NK0?!xeUi7w5I!Y~lXqAsdgfFy~nx#5Gu!oHZzw&nP%!5{8&*KxY2=gS~#r8d8 za?s^X{egp6ya+MAc6t5I98hCzA+Spw^|<du!H`|W zKdz6e42^E5bujRj!PF$jea7}&)coQ3>0e_81kK*9q0~VA;&1cEd?AhjN@^G{GfQD8 z?n7|tF^3+U_3@dD;lXrfd!j>^vL6YQzZ&B6XO&{1>hmd9H$|nC%b0$PU(Z2&Oc1z5j^N~o+Y)mGqW>GR>@Q2zr@NMB)^&~A5OR}JRv%=G4y;B zR<9!CzFG0gQ7&VOL%fWBK!>nt%hqG7b zv{hi{T-6r1;e1*1=93&4P4APQ*Fl29A$;pDQmMk1o~YCE&2vw{Z=RjgwiP6r<1IOh zBcEpm?|Cgz(d!yqJu=+3NvL+yY~l!;t`jc#^JOo%CCCJ-eaN3mwYfLSBum(;4%(jIaQne#bM4(lVVk*hlEaRiB` z+Yqkzr$tLS`q$aRfvg*}mt$Z=@Va#AIzaFxdnduId*KFK*qO5P2#tSgO(l}2HMRjD+ z_*GkOec7oYBohP<&$BJC1`jozSCb6MCWbGnZ6$G6( zU8o0!d|Fd^1%zw`m2HVSmVfB94`e@aMV7vv-V9Jfo5Wh>05GYuh6KXH_859`B+{Q_ zx7BlU^tAe&wsu~y+RC-rS8$bk^)*-R{vPjKc8+?$HH9ulL7FyRsf}Bgh%*9MHG|hw z@!;hjb3)Gh1rzM+@-|E$KL8BPmP1TaFOsNL@ zsKjt_XmMO!eR;@QTkdh>`RJa$Gxvb&5(%lGiZbzt=sF3heBn?omOpXizHp5Qw?`uP zoJc?X%N;fv5JC=F=YBRKoSid3IE$KG)VC?HjwZ~~!%%I@JEgPtv#|b*c%N_mdl`CcosgB87~$-;W@lm1n#;}ItZ9pb*8fEp+k6fF<;?sF){wV6C) zSxl9xUNp{~nbb~GT|ID-)(qp*ljQL1+N=vns!Jki0{ab2j~dN1WTk&JH-_j080bD5 zDW{R^xoGMl#8uBVEOYlqIk}_lp<&AB(-^ejwjEF1uhy zjDNg=mrAO{bh!=m<;xFhX$NSC!*VY`k*+B`2izUL>JoRk3`J`WS{Q3B_|Y=h>cK@H z18ylw;0`?bd!3X9qaU~s4D$RPK}Dxyo1`jxK|g)4VhW2dXoX`(rri(SstmjfQNxAd zIvj!s^CEDX8Y;h{5c|Zr=gi}cv}UMUhX#$uq3_r%kQla@WwHqucVbML07F_$W2`YK zarlz(`P!l0XjBDAOB&%|t6>|0a+U5*#%$APoMJTdkdEn{poXJrMt>I&-8^Nv0Q>g& z<1HI&Yp3vu+X4FGLqG(h$V2%g409>EX@ImCkHiVw=}ZZnxBU)cT?}) z_@|)yH!K<8WMMG_SooBb6nSL2n)o$xW%O_RSoL^wQ;WFbY(yS4QhsrGp4lcEtn+wM zbv1RP{u;Q2L?^93Pm~x>mMVFu;~`$Y@R{1$~^T)|bl9;&FL8{o4Des2&K}diDRj{t40Za`#CNP;d}pai(h`yP+7B00iGw%+1>lS$6x5=Jzln?g%y!`O z=Sa)=iNDOQC(|2F;pRe)Lp9YAbMWMKCw7lA%0hKl76Wft!OgK(k2P)ykOl)N2r+9g zmx05c0Wpq$bL!bV{tyVy$|~{dH5%gairJc|%MI|FFe(+yJUqUZ;B;<#IsL4u=*X#3 z%b77teht#WzG&?WEQM#K38d3nv@Tgb3$|y!E}StNVN|^_1j*8Qtdl-$;n((ur(pku z|Ep9U12gIe5}H}6)m1eq^5RSq?|*rSug`x%y0hsjN3)BQT_7Tw-a3pg6m{pExF;%Pjhk2kDzg@iI^7;}bEW5qPaeB*_ zVbC2}3R+oH_Zms@^Vk32ZtMpiCM1lP;)qf}XVDC{JtddT6DD0 zToX7*;xGRf%l)7CKF%z4B1?;l`_1|4H7@9;4+#AlJ9;d@i@xLMRDGhAPDk}JQDoW2 zDbu(&qo(dGbzl&XP$Ac~^jxj|=BOe~ZNfauOaN~zVwRj-XKURZ-AvcyXJy;$k~-4W zJ6yQ2YTJsUpagQ21GMA8wmUk*WhXzfQ7Z%g?5f4UdBvT4x`FxGtEk@vy3kM9kHt#zze9Y~pNI6X0CNY?pYn zk*zm~!u_8a*Vm)?M-4?IE<11r!%Qqqge=hj=+T2AhO43XmS;sk?1jfuvycM|cAW;- zHwVd?kq;IuJ@rQi+fcQ3r|axiM^G;FL*%1{x)uO-RX;jyYVW_BxlmMw8Dny#uNkC9B}n6i91#`dRr{a%KG-tnnl0u%rz-_*)pU?8hUxAt0WcC z_s*1#^*0nFWH91__k;fxO@aES7LZ!+_vH>Arwpe8scm5>d1tG5f)%~93&aDhiBij5 zCOXwSM|1M#+7Shu1DSeg(--j2^F850_1+UcG5 zjsvQq=ZrOST<1 z?tXs`=6&4y>X!bTDh-bJ7g_g%{U^z^b_{GCfN!n%A#!!~=Zh&p!c5|r@W7cGfUlAy z-g<_7jlEwTbZ>ARze@B%;;Bh%`0a3)`+5|0>UKJRFmnC!fxozs%d69s)Wt<(Z*Ol8 zxq$mk)p58Aw8YYAF|A@l0#_0#FmOlL3bL0c(7JeHKlnI0WWhH<@ zNvlw2Z5S1+h9$w&7ZogtQX!Rzgqus*SrKT=z_IiC7^)KqM5r}^)-SHXV7nigeBsj9<&+J|v!2Pda zZc>4;kofesR$xnw77?ToqwUZ2UjcD=Dhx8rj=!56o16^r_YeN>bkhr%?9WpHWSyZ+ zLqL{9olWLhJ2I<4Q)XC+#bRaA7KaQuU|HvyfzE)bI#K{~D0U{QH#vX_6>llmD(b9l*aNJBe7U$4%h$2rY!EJQ#9u zQaLMEKwWWW8G1?)D4|sxQGQYW{__^j=3)cv0E6ErB(8SdO+ShIhX*JB+lae5JA(rQ z?Oj|nK}cIY+nce(Mse~7#Q=R56W1i1cVI%{##=vf>}+FW4LPvh+FXMGxd`t13X16d zHQNvO{tmn*CMJFX0iwdfFOc#cQ~~!hs)1zR@QV3A|_Nqb-^S=K=qsM>*6=)?G%2RLa>%+9b**7=0)?Lht_M|dT zYqky|L3CEy)jJw89DN+my?M;6Gt%g9FJ1F`1QaWci&hxpCv?8YY)-!z_ozrH6F0+I zh?X4VyG2yT3JZK{N=d1lR8Lk?Hk-~%Bk0@RhOag=2uiPoVjS1adpn+vv)nhM_tse6 zm4qp$whMdg?N&7sUv~?`1s`6eF_NX2IQWZwObB)Cb#JW;|GxX1>KL~L z+2tN+JiNZl)^w)Y^vGPPCdF~zT{)qcRd}=~HF&h|&CzV*pz~En>TgOVgEMxo&Hes{ln51dH$w}rW*%MF{dZ;Lk()OD_wd0(>4NnNRl4?=^%JP$y z`{O$JhAqn#S8j$>p(+RoNgipKG3bnb0VISJ|MGr-I<_xWau!GvIXN+aGNekg2NE70 zo*ZyZml|()O>~VGP=W6{hi8%ZIPck7??0S{&=gRU1vlUbJ@&QehQlh|gsQp}#Au)UF=f~1# z$N1SAI9?B*4rB~<^|izi`L2z9LSw!=J|U?g>9P?F>|S0yw+IOOl11SbH|CIJ752ON zORP_&n!?cc!vfz?8@TrGhd}7DCJW2Wna#eLw~T^l$rl7$+ib%p9p10h8=dYbR+FFH0IF`dF5Vk@EF8HIBqOjE@^J8!D?GB$1(vF?8@rq?DYH=$*&lu9gx zQ*3I=TwUKZh#=>gW>PvnYxErdY@jcoX5ZlT* z8$g=@pbBmIPx&bBXS0Wil&Pud>(_yGb)Ep#c7BY9%OD&z7}_Bjdtdh^$n%3HeTFe%8+HF7dIS_ zi&SX4A1At4q3z?}&BzVwA4}VpJ0fBbvOoKAWAF}Xf`IYKP~=S(;#22$cnRtABl-yp z{(F7houl>dS<)NjXcDZ|;NH2W)GXHaZH`y+H=2*F`5G8aEaXFumKzHdi`nCqi=?Ea zP5Lckh|v*Dw63%Nb8G+GjUS6&;V^q1W}VUD;bGvq;tsMlLm?rdyNMrlxt&gZabvV4 zR9g1Z3&RpZHt~nsq`7@iaD2O5o-#M5bc7DP>xxhGMVk2pbBLx=-wc?yB$n@AHPFtpC4=i=vD}_0D zPY=InbgLb$g(#?mSC)8z7f(75QLk+>L0NOhrbHME{7fB2AIBk~{42Ws>c`Mb688W) z6RHaf3jtBrIv`=x!@!*!H^t6!zAL;1%0YUmLrBTm(5~0v3WDW;rJ)4Grs($HJ9al| zN%sI&M;%h1XdEX#6-s`kTATf{Owt^Hn8V=M7H9;a^uX^p8kcA6;<*gb3lWk>_}STmX892R`}8z)bYcs--GJz? zKV9ZQmaePo37{$UW7NFZr;>K1VNptW=HZI#|f+z!JRw7 zg3*Ypt1GV!df)4DUE?w?eska7`~I%?=kvX{R|M+l9)*6Y<+|L`on0c1 z{fF#M-NU`9Mxk-;oS(HeyTw+sH)>9QWOLUs@b57MI79mE9N;w)Ky_dC;_cWhmFpXO zLta|OptJPdBOC1-$NM_qxG(gzS3{&f6xmuZVxpmmj}9xxAULebj3}d*%~m$n3l{zate*IHlRsnGuJ1q*ZhmWwq69m-wRn z1!~y3;B_(gaG6$$%<<9bnd#}7;jdo~6-7Ve8V+o<;Y}rw1xwUw#qG+xqtM-FD$|R< zRJ2r`^{{)jxn~)pRV>Jtc!E0+ZDS&;Zo1bt&Q1aOULsWfeEMcthn00#HKZgZCEJfh z69`BsdaM`uSxATqy6SibAi#J5kJW#Eb-r0;{V$t$nst%s^f}+UDLhGt)=gRh}C^@Tm}4+x^j#b#^wlE_`uSQ zXVY^ZbADk9Pr23qekR8Kl7A7sc=H>2ImvO4Co;{78z35d{YkZsw#VmJWjE?Qsns6D zElY^5Oa(%7HGk4QBjzM_wf626S63wT-W|F|@tt_`zd2Qgg4F`&#z?@S!e^#a-9{4M z9u~WBs4()VB$P2ll1ybjC=gM~G0iS&{suOxUxegcd}7Y6D1@HLzF2)GyE9#rte{%g zYH;M!;z=T2gfQXn1EC!$*?n?s%nERVj!ZVq-5@?(LnE1)MmcgRl^S|W2?{5_;N(m^t>u$FE<}FP1IR~F7Gq` zK-AS%Q}Yqo{*<>fAVv-E$jv0r~=(nVLq1hug4BvcMhG+!fk1S8rQ|D`VGlx1KaM-jooB z9YNfIXjVoKdx0V+&Rm>mmsGYZ4dNDK4GT}wD(bnExQ7oPwzQm^%fHxRYv^cejji95 zBJ5}+p^|^_9CEs^&(4u3slxkfdRUe3r*muX|L+3w|C$Ya|9?q0`1kDo>76;d9Y~*5 zEkVow$(jg!mcY7*yUq5F?6a$yIXABT?jJ}i>_24m`0Ar=s^+doE8T@9IrhyuhpKdO zI#VBfWfE5f;$4rw2x|VvKje^@$$=jmaaOF(gELQ(ODj*UJ^VYGWo?m?jnJO-2trvK zKe#lo#!)b2E2xGSD>qxK=-env`x7+1U8hO?~>1ZqH8Z3EaX;jw%WF$!yw4TzW%^iUJj*pFr zbg7~jMwTdz(TL8@seCPYfl8BDV#R1@03$ zE#!r=f8^^JL{lhzt||p)1L5d^tj=@><`JDw?F%sgF)`SI^U(>-0$iyQcY+ABEJ*I} z$9vGT*3Z;kjHQR65XhGPS;M8U@E2Q2;(8Ci96+!V5>fB;Tk+0CE@ zAt50pG=-+)qFE04`w$X1S7qS30Q+C;Ms33W`S#9wfna$?vGmQ+V}plnHi~UWIVNtG zXu3T8X@lLV8#TYJyVPTXExAd@`As~Vhc*JS(K{eE*#&7LGB z0f%%J!n9w~s9jxaj~>E$H-gyfy*fU}qkv;dzvboTHrc8@fcspLy~pU=7X7_l=~k+v z;IoG_NpdmKoV`^D)0X9~x9yRptj8Og)tvbdeY^|5&{Xa<7F^1hlkDWa`#%St6P=-!U2=1C^Co-e0!0K^#+sU%jK+?prl#4rkf&`& z#uL45z(qfJ@SsO|>d@Myn`>W*R(W3TT35VwAKt6t(V0f@BC|{711Ka%%uLcXIli`= z=eWaULVTN(k{rEWPbNcw#+`1V1|1wD7Ws6g?h_pl3c(*`^SwGT3w10D7(X161g*|B z7TD^u`fkD~AnDqgrd<9t1)6I-ClCx*zC0az|C)g2^th8l?{`qCm9w>*VA3`0%LIEZ z2%a5siLp{yyLG%Bu3ei*(hr2|8yX|`NaKf=h9dwTgT0}vs;L7Fol`-Nh^Y0GKr#d- z(CQ#3tKb9ZlP}I@gV;Av&7ek)P29q8IikSl$(5Wcol?(8tr?v{rw?Ix;$y;cZ+0DP zw=(Tg9ILYvvtydLi_~s?_PWJN?YLnj4(B<(N>0}AwA{7XQvtJB+k@b9VgqaXq?)=xI zD)k+`o)t>`OZ1+e9@wg3+Q4j*1-!bhYEe#(Cv-&^+6(gQFLLjC6=fjB?}jL{7$YY@1SAjpfnns9fo>=6J>2BTA>z`Vs7dWC+lguROq5LQs))%ZQXGTT4Tt^R%<9RKg z8Xok7@t6`(@$m!Bs2cx-CP_^#SzWqikuNtkW2Ij(7Au7vC*IB}9e&`w9BnKEGgCNp z=>7Zm8;zmG)7dj;?p}r>JCC!p>D1BNGrJQzX*8OTf~w{PV3m$M(hjrG?(;rz{(Ngx zeaHO#y!T}|VXO-bn*UR7sHvdMKo=)p>>U1G@uiGXM7^$;xjUf=@|oofokf$Ih1VvQ zF=nmIISv8Bwty5&n5yQ(S4OW3cHrauoo?`tBtLI8pj-N7P0Z%+F*8N5ZLFF3J_4q5 zb(*-v!u&ksZQDK=EPgz3f6E@|o-Pd=K*y(5vx~x_2S%TEfgsVo+*kONa@m(B5RQii zn#q%9=#|`=3g(umO3q~F*|TSVZRmiDAP1vsnZm3|>J*d~w)}!nP^guDotoO`)2x(d ztr&d${?CRjNRnMA^%mg)`TV`>BmB-=!_WZC|l%=Hxc6QKP>tO~^ec?7rXc|y+ zx-CiHwI$Zv%1Y}>ORRV6i(?@M0%*m_%xHiM*vk=msl-^6jm-+pIm*wCLFz&Wday;V z+uQepF4Wb7#tq~;o>--A+57DP_3ea&d$`g-jArGFaduzF;TOT^s;|{ucJt@VnM!zc zex{nXj77*mpD6G4Zj?-TC&7wrKU@D%*q_)qPEuT%q_6j@uQ3;m?op`<3|i&{+Z-_g z=_AQ#eV6x-*M_c?P0pnIwx1swtYpzE;9i(xA;0*!i6(-}5CBiJt`hzn4%q^zmyn^Z ze41;pa~=(PMg#&GH)AQH3No-Bs=`hj*zCHlOBKEoYA|hf-k#$>Y;2Cf#RX7O)21m7 z#z-JcicVQiS}6G67PqEZLl1U8K7TlIFmQ! zJPEp+b}mj#h6YL*I8&BZRmH!*rmf0Ak=ecZK+$&DvcbEXz#I%m-5C3^(+_rZFgFQg zCtEaw1iphj(wOsx$ukD@6p zq$2^tzm|fz^?o#S#BP=|)?$+Rxsf>R8VnniL?VG54HN8N>bu#-fR!%&+rE9ypq+D- zVC#c<93`9lFWDvZc;LSR`!qyzXs;w1JCsSTia$)6ffQ=+?VTY9-Ougql1eCX{6lrK z?;=+Ja=MN`AegI?Ul~8wEa;vm3T(ifNVzL#qD#I;VDw5z%#Kx=k((Bcqx;XJsD7Ua zg5vDqtxM7I)@K!IwB7DBo3i;5+xun+9c_f1{6l{{{+wu*U71=-$_?;n0QSww%#2`q z$dPWjF*5!fh^R)!KYxCU%P5fr&kHim$jE@!Ki6|$Rm-W~ z1WaXvCSrSPYiq%$vG);o{qSK&&bKxi^s-P>6=@l5;!AuUk~KCRH$V2=k?h%Iy?S{p zpJ5g}*DrlPlHnPtSzP!oMX7aw>e&eNVEC&k|GEM=-sn(ef>Jx!Z?yMF7vk(UwR zvcOBZRg~~zIVYe1Bn@U7@Rzx)`>L9yOMnBYiylN*++sZ7;h}au$Fd4}rRCXIid&hN zOTt3gQosK0@vYU&{@A18@mKwi?9tDd_43f-m%#>Ar|UK{Y$YTgw_U?mTY_ zl&5lY5ATj*t^uNYE7!R0Pz|l^E*S{c-4Vg)Em5)VyoGcCYJI_#R#d(VI9!wcV#3db z0z)DO+q`SQR3|_zXK%0XMW;B-{gpWHRCY&TQzZ~q=RB4c+s&ZaO#ld@=1-4mJ=~S= z(cH&%PD-D~u3C81@rOfXtB~_wrEU2^HIgS>$<+a;9Ld6E*I>`!vtBmfhjctNwsW)Wiaad4)9GqVSGEMPk)X<4)%Lfy|qLi4PNle2z%nkF!P)~ zS*IT9^lhrXYN7a?tzNde|NG;A|M}lP{i|ibcRI(zc*hLNHb3C}7bd{O$n<2+afiEq E2S^E29RL6T literal 0 HcmV?d00001 diff --git a/media/docs/pythonDSL/cute_dsl_general/dsl_introduction.rst b/media/docs/pythonDSL/cute_dsl_general/dsl_introduction.rst index af487563..73b7e236 100644 --- a/media/docs/pythonDSL/cute_dsl_general/dsl_introduction.rst +++ b/media/docs/pythonDSL/cute_dsl_general/dsl_introduction.rst @@ -5,20 +5,30 @@ Introduction -====================== +============ Overview -------- -|DSL| is a Python-based domain-specific language (DSL) designed for |DC| of numeric and GPU-oriented code. Its primary goals are: +|DSL| is a Python-based domain-specific language (DSL) designed for |DC| of +high-performance GPU kernels. It evolved from the C++ CUTLASS library and is +now available as a decorator-based DSL. -- **Consistent with CuTe C++**, allowing users to express GPU kernels with full control of the hardware. +Its primary goals are: + +- **Zero-cost abstraction**, DSL is a zero-cost abstraction thanks to Hybrid DSL approach. +- **Consistent with CuTe C++**, allowing users to express GPU kernels with full + control of the hardware. - **JIT compilation** for both host and GPU execution. -- `DLPack `_ **integration**, enabling seamless interop with frameworks (e.g., PyTorch, JAX). -- **JIT caching**, so that repeated calls to the same function benefit from cached |IR| modules. -- **Native types and type inference** to reduce boilerplate and improve performance. -- **Optional lower-level control**, offering direct access to GPU backends or specialized |IR| dialects. +- `DLPack `_ **integration**, enabling seamless + interop with frameworks (e.g., PyTorch, JAX). +- **JIT caching**, so that repeated calls to the same function benefit from + cached |IR| modules. +- **Native types and type inference** to reduce boilerplate and improve + performance. +- **Optional lower-level control**, offering direct access to GPU backends or + specialized |IR| dialects. Decorators ---------- diff --git a/media/docs/pythonDSL/cute_dsl_general/resources.rst b/media/docs/pythonDSL/cute_dsl_general/resources.rst new file mode 100644 index 00000000..0a63a87a --- /dev/null +++ b/media/docs/pythonDSL/cute_dsl_general/resources.rst @@ -0,0 +1,28 @@ +.. _talks_and_presentations: +.. |DSL| replace:: CuTe DSL + +Talks and Presentations +======================= + +This page collects talks, presentations, and other resources related to |DSL| +and CUTLASS Python infrastructure. + +Conference Talks +---------------- + +**CuTeDSL: CUTLASS Python DSL Infrastructure** — *LLVM 2025* + +An introduction to the |DSL| architecture, covering the hybrid AST-rewrite and +tracing approach, MLIR code generation, and integration with CUTLASS. + +* `LLVM Video `_ +* `Slides (PDF) `_ + +---- + +**Enable Tensor Core Programming in Python with CUTLASS 4.0** — *GTC 2025* + +Learn how to leverage Tensor Cores directly from Python using CUTLASS 4.0's +new DSL front-end, enabling rapid kernel development without writing CUDA C++. + +* `GTC Video `_ diff --git a/media/docs/pythonDSL/overview.rst b/media/docs/pythonDSL/overview.rst index c98bea59..fbd3abd8 100644 --- a/media/docs/pythonDSL/overview.rst +++ b/media/docs/pythonDSL/overview.rst @@ -105,4 +105,4 @@ You can: - Propose support for additional data types or kernel variants - Help prioritize roadmap features by upvoting GitHub issues -Thank you for helping shape the future of CUTLASS DSLs! +Thank you for helping shape the future of CUTLASS DSLs! \ No newline at end of file diff --git a/python/CuTeDSL/cutlass/base_dsl/ast_helpers.py b/python/CuTeDSL/cutlass/base_dsl/ast_helpers.py index fa3c1e8e..cda5137d 100644 --- a/python/CuTeDSL/cutlass/base_dsl/ast_helpers.py +++ b/python/CuTeDSL/cutlass/base_dsl/ast_helpers.py @@ -179,12 +179,12 @@ class Executor: def ifexp_execute( self, pred, - generator_targets: tuple, + block_args: tuple, then_block: Callable, else_block: Callable, ): assert self._ifexp_dynamic, "Functions must be set before execution." - return self._ifexp_dynamic(pred, generator_targets, then_block, else_block) + return self._ifexp_dynamic(pred, block_args, then_block, else_block) # ============================================================================= @@ -309,16 +309,14 @@ def if_executor( def ifExp_executor( *, pred, - generator_targets: tuple, + block_args: tuple, then_block: Callable, else_block: Callable, ): if not executor._is_dynamic_expression(pred): - return ( - then_block(*generator_targets) if pred else else_block(*generator_targets) - ) + return then_block(*block_args) if pred else else_block(*block_args) else: - return executor.ifexp_execute(pred, generator_targets, then_block, else_block) + return executor.ifexp_execute(pred, block_args, then_block, else_block) # ============================================================================= diff --git a/python/CuTeDSL/cutlass/base_dsl/ast_preprocessor.py b/python/CuTeDSL/cutlass/base_dsl/ast_preprocessor.py index 52c82144..5eea3018 100644 --- a/python/CuTeDSL/cutlass/base_dsl/ast_preprocessor.py +++ b/python/CuTeDSL/cutlass/base_dsl/ast_preprocessor.py @@ -45,6 +45,7 @@ from typing import List, Set, Dict, Any, Callable, Optional from types import ModuleType from collections import OrderedDict from copy import deepcopy +from itertools import chain from .common import * from .utils.logger import log @@ -302,6 +303,7 @@ class SessionData: import_top_module: bool = False region_stack: list[Region] = field(default_factory=list) generator_targets: list[str] = field(default_factory=list) + lambda_args: list[str] = field(default_factory=list) @contextlib.contextmanager def set_current_class_name(self, class_name: str): @@ -2063,6 +2065,19 @@ class DSLPreprocessor(ast.NodeTransformer): return self._visit_Comprehension(node, key_value_visitor) + def visit_Lambda(self, node): + current_lambda_args = len(self.session_data.lambda_args) + for arg in node.args.args: + self.session_data.lambda_args.append(arg.arg) + + node.body = self.visit(node.body) + + self.session_data.lambda_args = self.session_data.lambda_args[ + :current_lambda_args + ] + + return node + def visit_ListComp(self, node): return self._visit_Comprehension( node, lambda n: setattr(n, "elt", self.visit(n.elt)) @@ -2113,7 +2128,10 @@ class DSLPreprocessor(ast.NodeTransformer): posonlyargs=[], args=[ ast.arg(arg=target, annotation=None) - for target in self.session_data.generator_targets + for target in chain( + self.session_data.generator_targets, + self.session_data.lambda_args, + ) ], kwonlyargs=[], kw_defaults=[], @@ -2129,7 +2147,10 @@ class DSLPreprocessor(ast.NodeTransformer): posonlyargs=[], args=[ ast.arg(arg=target, annotation=None) - for target in self.session_data.generator_targets + for target in chain( + self.session_data.generator_targets, + self.session_data.lambda_args, + ) ], kwonlyargs=[], kw_defaults=[], @@ -2151,11 +2172,14 @@ class DSLPreprocessor(ast.NodeTransformer): keywords=[ ast.keyword(arg="pred", value=self.visit(node.test)), ast.keyword( - arg="generator_targets", + arg="block_args", value=ast.Tuple( elts=[ ast.Name(id=name, ctx=ast.Load()) - for name in self.session_data.generator_targets + for name in chain( + self.session_data.generator_targets, + self.session_data.lambda_args, + ) ], ctx=ast.Load(), ), diff --git a/python/CuTeDSL/cutlass/base_dsl/dsl.py b/python/CuTeDSL/cutlass/base_dsl/dsl.py index ef17f185..0305ba30 100644 --- a/python/CuTeDSL/cutlass/base_dsl/dsl.py +++ b/python/CuTeDSL/cutlass/base_dsl/dsl.py @@ -50,7 +50,11 @@ from .jit_executor import JitCompiledFunction, JitFunctionArtifacts from .utils.timer import timer from .utils.logger import log from .utils.stacktrace import filter_exception, walk_to_top_module, filter_stackframe -from .runtime.jit_arg_adapters import is_argument_constexpr, JitArgAdapterRegistry +from .runtime.jit_arg_adapters import ( + is_argument_constexpr, + is_arg_spec_constexpr, + JitArgAdapterRegistry, +) from .ast_preprocessor import DSLPreprocessor from .common import * @@ -1310,8 +1314,7 @@ class BaseDSL(metaclass=DSLSingletonMeta): dynamic_args = [] dynamic_kwargs = OrderedDict() for i, arg in enumerate(args): - if not is_argument_constexpr( - arg, + if not is_arg_spec_constexpr( args_spec.annotations.get(args_spec.args[i], None), args_spec.args[i], i, @@ -1319,7 +1322,7 @@ class BaseDSL(metaclass=DSLSingletonMeta): ): dynamic_args.append(arg) for i, (k, v) in enumerate(kwargs.items()): - if not is_argument_constexpr(v, args_spec.kwonlyargs[i], k, i, funcBody): + if not is_arg_spec_constexpr(args_spec.kwonlyargs[i], k, i, funcBody): dynamic_kwargs[k] = v return dynamic_args, dynamic_kwargs diff --git a/python/CuTeDSL/cutlass/base_dsl/tvm_ffi_builder/mlir_builder.py b/python/CuTeDSL/cutlass/base_dsl/tvm_ffi_builder/mlir_builder.py index c144bcd2..9fc1e02b 100644 --- a/python/CuTeDSL/cutlass/base_dsl/tvm_ffi_builder/mlir_builder.py +++ b/python/CuTeDSL/cutlass/base_dsl/tvm_ffi_builder/mlir_builder.py @@ -371,63 +371,6 @@ class MLIRBuilder(MLIRTypeBuilder): self.const_str_table[content] = symbol return symbol - def get_or_load_global_func_ptr_from_text( - self, - current_block: ir.Block, - function_name: str, - ) -> ir.Value: - """Get or create a function pointer global in .text section and load it. - - This creates a constant global function pointer in the .text section - (for AArch64 ADRP range compatibility) and performs a volatile load - to prevent optimization. - - This forces the function pointer to be local to the code, bypassing GOT entry - ADRP lookup issues on AArch64 when GOT and .text section are more than 4GB - apart which can happen when ASLR is applied. - """ - # Check if we've already created this global - if function_name not in self.const_func_ptr_table: - symbol = f"__func_ptr_{function_name}" - - module_body = self.module.body - with ir.InsertionPoint(module_body): - # 1. Create the global constant - # We use 'private' linkage so it doesn't conflict across modules - global_ptr = llvm.GlobalOp( - self.ptr_type, - symbol, - ir.Attribute.parse("#llvm.linkage"), - # Initialization via block below - ) - - # 2. Set the necessary attributes for JIT safety and AArch64 range - # We use 'constant' to mark it as immutable - # We use 'section = ".text"' to force it into the code block - global_ptr.attributes["constant"] = ir.UnitAttr.get() - global_ptr.attributes["section"] = ir.StringAttr.get(".text") - - # 3. Add a constructor block to the GlobalOp to initialize it - # with the address of the target function - initializer_block = global_ptr.initializer.blocks.append() - with ir.InsertionPoint(initializer_block): - # Get the address of the external function - func_addr = llvm.AddressOfOp(self.ptr_type, function_name).res - # Return the address as the initial value of the global - llvm.return_(arg=func_addr) - - self.const_func_ptr_table[function_name] = symbol - else: - symbol = self.const_func_ptr_table[function_name] - - # Load it with volatile semantics in the current block - with ir.InsertionPoint(current_block): - symbol_addr = self.address_of(symbol, self.ptr_type) - # Perform a volatile load to prevent optimization - load_op = llvm.load(self.ptr_type, symbol_addr) - # Set volatile attribute to prevent optimization - load_op.owner.attributes["volatile_"] = ir.UnitAttr.get() - return load_op # function def function( diff --git a/python/CuTeDSL/cutlass/cute/algorithm.py b/python/CuTeDSL/cutlass/cute/algorithm.py index b4d96f12..bc03231c 100644 --- a/python/CuTeDSL/cutlass/cute/algorithm.py +++ b/python/CuTeDSL/cutlass/cute/algorithm.py @@ -29,8 +29,14 @@ from .core import ( append_ones, group_modes, ) -from .atom import MmaAtom, CopyAtom, make_atom - +from .atom import ( + MmaAtom, + CopyAtom, + make_atom, + _normalize_variadic_tensor_operand, + copy_atom_call, +) +from .nvgpu.common import CacheEvictionPriority def _normalize_gemm_operand_list( x: Union["Tensor", List["Tensor"], Tuple["Tensor", ...]], name: str @@ -156,7 +162,7 @@ def basic_copy(src: Tensor, dst: Tensor, *, loc=None, ip=None) -> None: src.element_type.mlir_type, src.element_type.width ) simt_copy = make_atom(simt_copy_ty, loc=loc, ip=ip) - return _cute_ir.copy(simt_copy, src.value, dst.value, loc=loc, ip=ip) + return _cute_ir.copy(simt_copy, [src.value], [dst.value], loc=loc, ip=ip) s = size(dst, loc=loc, ip=ip) # Always generate an scf.for Op when one of the tensors is dynamic @@ -210,7 +216,14 @@ def _basic_copy_if_static( @dsl_user_op -def autovec_copy(src: Tensor, dst: Tensor, *, loc=None, ip=None) -> None: +def autovec_copy( + src: Tensor, + dst: Tensor, + *, + l1c_evict_priority: CacheEvictionPriority = CacheEvictionPriority.EVICT_NORMAL, + loc=None, + ip=None, +) -> None: """ Auto-vectorization SIMT copy policy. @@ -263,11 +276,15 @@ def autovec_copy(src: Tensor, dst: Tensor, *, loc=None, ip=None) -> None: # Dispatch to copy with atom simt_type = _cute_nvgpu_ir.CopyAtomSIMTSyncCopyType.get( - src.element_type.mlir_type, num_bits_per_copy + src.element_type.mlir_type, + num_bits_per_copy, + 0, + 0, + l1c_evict_priority._to_ir(), ) simt_copy = make_atom(simt_type, loc=loc, ip=ip) return _cute_ir.copy( - simt_copy, tiled_src.value, tiled_dst.value, loc=loc, ip=ip + simt_copy, [tiled_src.value], [tiled_dst.value], loc=loc, ip=ip ) # Failed to vectorize, use a basic copy @@ -331,8 +348,8 @@ def _parse_auto_multicast_args( @dsl_user_op def copy( atom: CopyAtom, - src: Tensor, - dst: Tensor, + src: Union[Tensor, List[Tensor], Tuple[Tensor, ...]], + dst: Union[Tensor, List[Tensor], Tuple[Tensor, ...]], *, pred: Optional[Tensor] = None, loc=None, @@ -343,10 +360,10 @@ def copy( :param atom: Copy atom specifying the transfer operation :type atom: CopyAtom - :param src: Source tensor with layout profile ``(V, Rest...)`` - :type src: Tensor - :param dst: Destination tensor with layout profile ``(V, Rest...)`` - :type dst: Tensor + :param src: Source tensor or list of tensors with layout profile ``(V, Rest...)`` + :type src: Union[Tensor, List[Tensor], Tuple[Tensor, ...]] + :param dst: Destination tensor or list of tensors with layout profile ``(V, Rest...)`` + :type dst: Union[Tensor, List[Tensor], Tuple[Tensor, ...]] :param pred: Optional predication tensor for conditional transfers, defaults to None :type pred: Optional[Tensor], optional :param loc: Source location information, defaults to None @@ -374,6 +391,12 @@ def copy( Source and destination tensors must be partitioned in accordance with the Copy Atom specifications. Post-partitioning, both tensors will exhibit a ``(V, Rest...)`` layout profile. + The operands `src` and `dst` are variadic, each containing a variable number of tensors: + + - For regular copy, `src` and `dst` contain single source and destination tensors respectively. + - For copy with auxiliary operands, `src` and `dst` contain the primary tensors followed by + their respective auxiliary tensors. + **Precondition:** The size of mode 1 must be equal for both source and destination tensors: ``size(src, mode=[1]) == size(dst, mode=[1])`` @@ -399,41 +422,54 @@ def copy( for future releases. """ - if isinstance(src.type, _cute_ir.MemRefType) and isinstance( - dst.type, _cute_ir.MemRefType + # Normalize src/dst to lists for variadic IR operands + src_list = _normalize_variadic_tensor_operand(src, "src") + dst_list = _normalize_variadic_tensor_operand(dst, "dst") + + # Validate primary tensors (first element) + src_primary = src_list[0] + dst_primary = dst_list[0] + + if isinstance(src_primary.type, _cute_ir.MemRefType) and isinstance( + dst_primary.type, _cute_ir.MemRefType ): - if src.element_type.width != dst.element_type.width: + if src_primary.element_type.width != dst_primary.element_type.width: raise TypeError( "`copy` currently only supports equal source and destination " "element type bit width" ) - if rank(src) != rank(dst): + if rank(src_primary) != rank(dst_primary): raise ValueError( "Expected source and destination tensors to have the same rank, " - f"but got {rank(src)} and {rank(dst)}" + f"but got {rank(src_primary)} and {rank(dst_primary)}" ) - # Canonicalize to at least rank-2 tensors - src = group_modes(append_ones(src, up_to_rank=2), 1) - dst = group_modes(append_ones(dst, up_to_rank=2), 1) + # Canonicalize all tensors to at least rank-2 + src_list = [group_modes(append_ones(t, up_to_rank=2), 1) for t in src_list] + dst_list = [group_modes(append_ones(t, up_to_rank=2), 1) for t in dst_list] if pred is not None: pred = group_modes(append_ones(pred, up_to_rank=2), 1) - if is_static(src.shape[1]) and is_static(dst.shape[1]): - if size(src, mode=[1]) != size(dst, mode=[1]): + # Recompute primary references after canonicalization + src_primary = src_list[0] + dst_primary = dst_list[0] + + if is_static(src_primary.shape[1]) and is_static(dst_primary.shape[1]): + if size(src_primary, mode=[1]) != size(dst_primary, mode=[1]): raise ValueError( "Expected source and destination tensors to have the same size in mode-1, " - f"but got {size(src, mode=[1])} and {size(dst, mode=[1])}" + f"but got {size(src_primary, mode=[1])} and {size(dst_primary, mode=[1])}" ) multicast_attr_pairs = _parse_auto_multicast_args(kwargs) value = atom._unpack(loc=loc, ip=ip, **kwargs) - if isinstance(pred, Tensor): - pred = pred.value + pred_value = pred.value if isinstance(pred, Tensor) else pred - op = _cute_ir.copy(value, src.value, dst.value, pred=pred, loc=loc, ip=ip) + src_vals = [t.value for t in src_list] + dst_vals = [t.value for t in dst_list] + op = _cute_ir.copy(value, src_vals, dst_vals, pred=pred_value, loc=loc, ip=ip) for name, attr in multicast_attr_pairs: op.attributes[name] = attr diff --git a/python/CuTeDSL/cutlass/cute/arch/nvvm_wrappers.py b/python/CuTeDSL/cutlass/cute/arch/nvvm_wrappers.py index 329e2e49..ac109718 100644 --- a/python/CuTeDSL/cutlass/cute/arch/nvvm_wrappers.py +++ b/python/CuTeDSL/cutlass/cute/arch/nvvm_wrappers.py @@ -13,7 +13,7 @@ from functools import partial from typing import Any, Optional, Tuple, Union, Callable, Literal from typing_extensions import deprecated -from cutlass.cutlass_dsl import T, dsl_user_op +from cutlass.cutlass_dsl import T, dsl_user_op, target_version import cutlass.cutlass_dsl as cutlass_dsl @@ -2557,3 +2557,42 @@ def cvt_f4e2m1x8_to_f16x8(src_vec8, *, loc=None, ip=None): vec_f16x8_type = ir.VectorType.get([8], Float16.mlir_type, loc=loc) vec_f16x8 = llvm.bitcast(vec_f16x8_type, vec_f32x4, loc=loc, ip=ip) return vec_f16x8 + + +@dsl_user_op +def mapa(ptr, cta_rank_in_cluster=0, *, loc=None, ip=None): + """ + Map a pointer to distributed shared memory across cluster. + + Portable wrapper that uses the appropriate NVVM API based on CUDA version: + - CUDA 13.1+: Uses nvvm.mapa with dsmem address space + - CUDA 12.9: Uses nvvm.mapa_shared_cluster + + Args: + ptr: Pointer to shared memory (llvm_ptr attribute will be used) + cta_rank_in_cluster: CTA rank within the cluster (default 0) + + Returns: + Mapped LLVM pointer to shared memory + """ + if target_version(min_version="13.1"): + dsmem_ptr_ty = llvm.PointerType.get(7) # dsmem + smem_ptr_ty = llvm.PointerType.get(3) # smem + + llvm_ptr = nvvm.mapa( + dsmem_ptr_ty, + ptr.llvm_ptr, + Int32(cta_rank_in_cluster).ir_value(loc=loc, ip=ip), + loc=loc, + ip=ip, + ) + return llvm.addrspacecast(smem_ptr_ty, llvm_ptr, loc=loc, ip=ip) + else: + llvm_ptr = ptr.llvm_ptr + return nvvm.mapa_shared_cluster( + llvm_ptr.type, + llvm_ptr, + Int32(cta_rank_in_cluster).ir_value(loc=loc, ip=ip), + loc=loc, + ip=ip, + ) diff --git a/python/CuTeDSL/cutlass/cute/atom.py b/python/CuTeDSL/cutlass/cute/atom.py index 0ea53942..0d4ec36e 100644 --- a/python/CuTeDSL/cutlass/cute/atom.py +++ b/python/CuTeDSL/cutlass/cute/atom.py @@ -10,7 +10,7 @@ # is strictly prohibited. from abc import ABC, ABCMeta, abstractmethod -from typing import Type, Union, Optional, Any, overload +from typing import Type, Union, Optional, Any, List, Tuple, overload from .typing import Shape, Layout, Tile, Tensor, Numeric, Int32 from .core import ( @@ -1113,25 +1113,66 @@ def make_tiled_copy_C_atom(atom: CopyAtom, mma: TiledMma, *, loc=None, ip=None): return _make_tiled_copy(atom, layout_tv, tiler_mn, loc=loc, ip=ip) +def _normalize_variadic_tensor_operand( + x: Union["Tensor", List["Tensor"], Tuple["Tensor", ...]], name: str +) -> List["Tensor"]: + """Normalize a Tensor or sequence of Tensors to a list of Tensors. + + Helper function for operations with variadic operands. + """ + if isinstance(x, Tensor): + return [x] + if isinstance(x, (list, tuple)): + if len(x) == 0: + raise ValueError(f"`{name}` must contain at least one Tensor") + if not all(isinstance(t, Tensor) for t in x): + raise TypeError(f"All elements of `{name}` must be Tensor") + return list(x) # type: ignore + raise TypeError(f"`{name}` must be a Tensor or a sequence of Tensors") + + @dsl_user_op def copy_atom_call( atom: CopyAtom, - src: Tensor, - dst: Tensor, + src: Union[Tensor, List[Tensor], Tuple[Tensor, ...]], + dst: Union[Tensor, List[Tensor], Tuple[Tensor, ...]], *, pred: Optional[Tensor] = None, loc=None, ip=None, **kwargs, ) -> None: - """Executes a single copy atom operation between two tensors. + """ + Execute a single copy atom operation. + + The copy_atom_call operation executes a copy atom with the given operands. + Source and destination tensors have layout profile ``(V)``. + + The ``V-mode`` represents either: + + - A singular mode directly consumable by the provided Copy Atom + - A composite mode requiring recursive decomposition, structured as ``(V, Rest...)``, + + For src/dst layout like ``(V, Rest...)``, the layout profile of ``pred`` must match ``(Rest...)``. + + - Certain Atoms may require additional operation-specific keyword arguments. + - Current implementation limits ``V-mode`` rank to 2 or less. Support for higher ranks is planned + for future releases. + + Both ``src`` and ``dst`` operands are variadic, containing a variable number of tensors: + + - For regular copy, ``src`` and ``dst`` each contain a single tensor. + - For copy with auxiliary operands, they contain the main tensor followed by + auxiliary tensors. For example: :param atom: Copy atom specifying the transfer operation :type atom: CopyAtom - :param src: Source tensor with layout profile ``(V)`` - :type src: Tensor - :param dst: Destination tensor with layout profile ``(V)`` - :type dst: Tensor + :param src: Source tensor(s) with layout profile ``(V)``. Can be a single Tensor + or a list/tuple of Tensors for operations with auxiliary source operands. + :type src: Union[Tensor, List[Tensor], Tuple[Tensor, ...]] + :param dst: Destination tensor(s) with layout profile ``(V)``. Can be a single Tensor + or a list/tuple of Tensors for operations with auxiliary destination operands. + :type dst: Union[Tensor, List[Tensor], Tuple[Tensor, ...]] :param pred: Optional predication tensor for conditional transfers, defaults to None :type pred: Optional[Tensor], optional :param loc: Source location information, defaults to None @@ -1144,54 +1185,43 @@ def copy_atom_call( :return: None :rtype: None - The copy_atom_call operation executes a single copy atom with the given operands. - Source and destination tensors with layout profile like ``(V)``. - - The ``V-mode`` represents either: - - - A singular mode directly consumable by the provided Copy Atom - - A composite mode requiring recursive decomposition, structured as ``(V, Rest...)``, - - For src/dst layout like ``(V, Rest...)``, the layout profile of ``pred`` must match ``(Rest...)``. - **Examples**: .. code-block:: python - # Basic copy atom operation + # Regular copy atom operation cute.copy_atom_call(copy_atom, src, dst) # Predicated copy atom operation cute.copy_atom_call(copy_atom, src, dst, pred=pred) - .. note:: - - - Certain Atoms may require additional operation-specific keyword arguments. - - Current implementation limits ``V-mode`` rank to 2 or less. Support for higher ranks is planned - for future releases. - """ - if isinstance(src.type, _cute_ir.MemRefType) and isinstance( - dst.type, _cute_ir.MemRefType + # Normalize src/dst to lists for variadic IR operands, while keeping old API working. + src_list = _normalize_variadic_tensor_operand(src, "src") + dst_list = _normalize_variadic_tensor_operand(dst, "dst") + + # Validate first src/dst for element type width check + if isinstance(src_list[0].type, _cute_ir.MemRefType) and isinstance( + dst_list[0].type, _cute_ir.MemRefType ): - if src.element_type.width != dst.element_type.width: + if src_list[0].element_type.width != dst_list[0].element_type.width: raise TypeError( "`copy_atom_call` currently only supports equal source and destination " "element type bit width" ) - if rank(src, mode=[0]) > 2 or rank(dst, mode=[0]) > 2: + if rank(src_list[0], mode=[0]) > 2 or rank(dst_list[0], mode=[0]) > 2: raise NotImplementedError( "V-mode (mode-0) with rank > 2 is not supported yet, " - f"but got rank(src, mode=[0]) = {rank(src, mode=[0])} and rank(dst, mode=[0]) = {rank(dst, mode=[0])}" + f"but got rank(src, mode=[0]) = {rank(src_list[0], mode=[0])} and rank(dst, mode=[0]) = {rank(dst_list[0], mode=[0])}" ) value = atom._unpack(loc=loc, ip=ip, **kwargs) if isinstance(pred, Tensor): pred = pred.value - return _cute_ir.copy_atom_call( - value, src.value, dst.value, pred=pred, loc=loc, ip=ip - ) + src_vals = [t.value for t in src_list] + dst_vals = [t.value for t in dst_list] + return _cute_ir.copy_atom_call(value, src_vals, dst_vals, pred=pred, loc=loc, ip=ip) @dsl_user_op diff --git a/python/CuTeDSL/cutlass/cute/experimental/README.md b/python/CuTeDSL/cutlass/cute/experimental/README.md index 16679cc9..914330ab 100644 --- a/python/CuTeDSL/cutlass/cute/experimental/README.md +++ b/python/CuTeDSL/cutlass/cute/experimental/README.md @@ -51,4 +51,7 @@ - `get_cta_v_map_ab` — Compute CTA-V map for A/B operands - `get_cta_v_map_c` — Compute CTA-V map for C operand +- `make_tmem_layout_acc` — Derive TMEM accumulator buffer layout from a tiled MMA +- `make_tmem_layout_a` — Derive TMEM A-operand buffer layout from a tiled MMA +- `make_t2r_rmem_layout` — Derive per-thread RMEM buffer layout for the T2R epilogue copy diff --git a/python/CuTeDSL/cutlass/cute/experimental/utils.py b/python/CuTeDSL/cutlass/cute/experimental/utils.py index 5b230393..5a9c72be 100644 --- a/python/CuTeDSL/cutlass/cute/experimental/utils.py +++ b/python/CuTeDSL/cutlass/cute/experimental/utils.py @@ -70,10 +70,93 @@ def get_cta_v_map_c( :param gmem_tensor: Global-memory tensor being stored/loaded by TMA. :type gmem_tensor: cute.Tensor - :param epi_tile: Epilogue tile layout describing the CTA’s output tile shape. + :param epi_tile: Epilogue tile layout describing the CTA's output tile shape. :type epi_tile: cute.Layout :returns: A layout suitable to pass as `cta_v_map=...` to `tma_store` / `tma_load`. :rtype: cute.Layout """ ident = cute.core.make_identity_layout(gmem_tensor.shape, loc=loc, ip=ip) return cute.core.composition(ident, epi_tile, loc=loc, ip=ip) + + +def make_tmem_layout_acc( + tiled_mma, + mnk_tiler, + acc_stage, + *, + loc=None, + ip=None, +): + """Return TMEM accumulator buffer layout for a tiled MMA. + + This is a small helper around ``tiled_mma.make_fragment_C(...).layout`` to + keep example code fragment-free at the call site. + + :param tiled_mma: The MMA tiler (``cute.TiledMma``). + :type tiled_mma: cute.TiledMma + :param mnk_tiler: Full MNK tiler; only the MN components are used for C. + :type mnk_tiler: tuple + :param acc_stage: Accumulator pipeline stages. + :param loc: Optional location for DSL ops. + :param ip: Optional insertion point for DSL ops. + :return: Layout for the accumulator TMEM buffer. + :rtype: cute.Layout + """ + acc_shape = tiled_mma.partition_shape_C(mnk_tiler[:2], loc=loc, ip=ip) + acc_shape_staged = cute.append(acc_shape, acc_stage, loc=loc, ip=ip) + return tiled_mma.make_fragment_C(acc_shape_staged, loc=loc, ip=ip).layout + + +def make_tmem_layout_a( + tiled_mma, + mk_tiler, + stage, + *, + loc=None, + ip=None, +): + """Return TMEM A operand buffer layout for a tiled MMA. + + :param tiled_mma: The MMA tiler (``cute.TiledMma``). + :type tiled_mma: cute.TiledMma + :param mk_tiler: MK tiler used to shape the A operand. + :type mk_tiler: tuple + :param stage: Pipeline stages for the A operand buffer. + :param loc: Optional location for DSL ops. + :param ip: Optional insertion point for DSL ops. + :return: Layout for the A operand TMEM buffer. + :rtype: cute.Layout + """ + a_shape = tiled_mma.partition_shape_A(mk_tiler, loc=loc, ip=ip) + a_shape_staged = cute.append(a_shape, stage, loc=loc, ip=ip) + return tiled_mma.make_fragment_A(a_shape_staged, loc=loc, ip=ip).layout + + +def make_t2r_rmem_layout( + tiled_copy_t2r, + gC_mnl_epi, + tidx, + *, + loc=None, + ip=None, +): + """Return RMEM buffer layout for the T2R epilogue destination. + + Computes the per-thread RMEM buffer layout produced by a TMEM->RMEM copy + for a single epilogue iteration. + + :param tiled_copy_t2r: The TMEM->RMEM tiled copy op (``cute.TiledCopy``). + :type tiled_copy_t2r: cute.TiledCopy + :param gC_mnl_epi: Global C tensor partitioned by epilogue tile. + :type gC_mnl_epi: cute.Tensor + :param tidx: Thread index for the copy slice. + :param loc: Optional location for DSL ops. + :param ip: Optional insertion point for DSL ops. + :return: Layout for the RMEM buffer. + :rtype: cute.Layout + """ + thr_copy_t2r = tiled_copy_t2r.get_slice(tidx) + tTR_gC = thr_copy_t2r.partition_D(gC_mnl_epi, loc=loc, ip=ip) + return cute.make_fragment_like( + tTR_gC[(None, None, None, 0, 0)].layout, loc=loc, ip=ip + ) diff --git a/python/CuTeDSL/cutlass/cute/nvgpu/tcgen05/__init__.py b/python/CuTeDSL/cutlass/cute/nvgpu/tcgen05/__init__.py index f1afc177..74d1ce7f 100644 --- a/python/CuTeDSL/cutlass/cute/nvgpu/tcgen05/__init__.py +++ b/python/CuTeDSL/cutlass/cute/nvgpu/tcgen05/__init__.py @@ -19,6 +19,7 @@ __all__ = [ # copy.py # "Repetition", + "TmemLoadRedOp", "Pack", "Unpack", "Ld16x64bOp", diff --git a/python/CuTeDSL/cutlass/cute/nvgpu/tcgen05/copy.py b/python/CuTeDSL/cutlass/cute/nvgpu/tcgen05/copy.py index 76b2198b..0f165af6 100644 --- a/python/CuTeDSL/cutlass/cute/nvgpu/tcgen05/copy.py +++ b/python/CuTeDSL/cutlass/cute/nvgpu/tcgen05/copy.py @@ -26,6 +26,22 @@ from ...typing import Numeric from .mma import CtaGroup +class TmemLoadRedOp(enum.Enum): + """ + An enumeration for the possible reduce operations for TMEM load operations. + """ + + MAX = _cute_nvgpu_ir.TmemLoadRedOp.max + MAXABS = _cute_nvgpu_ir.TmemLoadRedOp.maxabs + MIN = _cute_nvgpu_ir.TmemLoadRedOp.min + MINABS = _cute_nvgpu_ir.TmemLoadRedOp.minabs + + def __str__(self) -> str: + return f"{self.__class__.__name__}.{self.name}" + + def __repr__(self) -> str: + return f"<{self.__class__.__name__}.{self.name}>" + class Repetition(enum.Enum): """ An enumeration for the number of repetitions of a given TMEM copy within the instruction. @@ -390,6 +406,97 @@ class Ld32x32bTrait(Trait): pass +@dataclass(frozen=True) +class LdRed16x32bx2Op(_LdBase): + """ + 16x32bx2 TMEM load Reduce Operation. + + See the `PTX documentation `__. + This Operation corresponds to the ``.red`` and ``.16x32bx2`` qualifiers. + """ + + redOp: TmemLoadRedOp = TmemLoadRedOp.MAX + nan: bool = False + half_split_off: int = 0 + + def _make_trait( + self, copy_internal_type: Type[Numeric], *, loc=None, ip=None, **kwargs + ) -> "LdRed16x32bx2Trait": + """ + Create a trait object for the 16x32bx2 TMEM load Reduce operation. + + :param copy_internal_type: The data type for the copy operation + :type copy_internal_type: Type[Numeric] + :param loc: MLIR location information for debugging, defaults to None + :type loc: optional + :param ip: MLIR insertion point for code generation, defaults to None + :type ip: optional + :param kwargs: Additional keyword arguments + :type kwargs: dict + :return: A trait object for this load operation + :rtype: LdRed16x32bx2Trait + """ + ty = _cute_nvgpu_ir.CopyAtomSM10xTmemLoadRedType.get( + copy_internal_type.mlir_type, + 16, + 32, + self.repeat.value, + self.redOp.value, + ir.UnitAttr.get() if self.nan else None, + ir.IntegerAttr.get(ir.IntegerType.get_signless(32), self.half_split_off), + ) + return LdRed16x32bx2Trait(make_atom(ty, loc=loc, ip=ip)) + + +class LdRed16x32bx2Trait(Trait): + pass + + +@dataclass(frozen=True) +class LdRed32x32bOp(_LdBase): + """ + 32x32b TMEM load Reduce Operation. + + See the `PTX documentation `__. + This Operation corresponds to the ``red`` and ``.32x32`` qualifiers. + """ + + redOp: TmemLoadRedOp = TmemLoadRedOp.MAX + nan: bool = False + + def _make_trait( + self, copy_internal_type: Type[Numeric], *, loc=None, ip=None, **kwargs + ) -> "LdRed32x32bTrait": + """ + Create a trait object for the 32x32b TMEM load Reduce operation. + + :param copy_internal_type: The data type for the copy operation + :type copy_internal_type: Type[Numeric] + :param loc: MLIR location information for debugging, defaults to None + :type loc: optional + :param ip: MLIR insertion point for code generation, defaults to None + :type ip: optional + :param kwargs: Additional keyword arguments + :type kwargs: dict + :return: A trait object for this load operation + :rtype: LdRed32x32bTrait + """ + ty = _cute_nvgpu_ir.CopyAtomSM10xTmemLoadRedType.get( + copy_internal_type.mlir_type, + 32, + 32, + self.repeat.value, + self.redOp.value, + ir.UnitAttr.get() if self.nan else None, + None, + ) + return LdRed32x32bTrait(make_atom(ty, loc=loc, ip=ip)) + + +class LdRed32x32bTrait(Trait): + pass + + @dataclass(frozen=True) class _StBase(CopyOp): """ diff --git a/python/CuTeDSL/cutlass/cute/nvgpu/warp/copy.py b/python/CuTeDSL/cutlass/cute/nvgpu/warp/copy.py index d284ea21..baff4839 100644 --- a/python/CuTeDSL/cutlass/cute/nvgpu/warp/copy.py +++ b/python/CuTeDSL/cutlass/cute/nvgpu/warp/copy.py @@ -103,15 +103,20 @@ class LdMatrix8x16x8bOp(BaseOp): self, "expects the 'num_matrices' Op parameter to be one of [1,2,4]", ) - if self.unpack_bits not in [4, 6]: - raise OpError(self, "Op unpack bits must be 4 or 6") + if self.unpack_bits not in [None, 4, 6]: + raise OpError(self, "Op unpack bits must be 4 or 6 or None") def _make_trait( self, copy_internal_type: Type[Numeric], *, loc=None, ip=None, **kwargs ) -> "LdMatrix8x16x8bTrait": - mode = _pack_shape((8, 16), loc=loc, ip=ip) - sz_pattern = _cute_nvgpu_ir.LdsmSzPattern.u4x16p64to8 - if self.unpack_bits == 6: + # LdMatrix8x16x8b without unpacking doesn't exist + # but is equivalent to LdMatrix8x8x16b + mode_n = 8 if self.unpack_bits is None else 16 + mode = _pack_shape((8, mode_n), loc=loc, ip=ip) + sz_pattern = _cute_nvgpu_ir.LdsmSzPattern.u16 + if self.unpack_bits == 4: + sz_pattern = _cute_nvgpu_ir.LdsmSzPattern.u4x16p64to8 + elif self.unpack_bits == 6: sz_pattern = _cute_nvgpu_ir.LdsmSzPattern.u6x16p32to8 ty = _cute_nvgpu_ir.CopyAtomLdsmType.get( copy_internal_type.mlir_type, diff --git a/python/CuTeDSL/cutlass/cute/runtime.py b/python/CuTeDSL/cutlass/cute/runtime.py index 87082656..6abbc527 100644 --- a/python/CuTeDSL/cutlass/cute/runtime.py +++ b/python/CuTeDSL/cutlass/cute/runtime.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: LicenseRef-NvidiaProprietary # # Use of this software is governed by the terms and conditions of the @@ -141,6 +141,11 @@ class _Tensor(Tensor): # If tensor is already a DLPack object, use it directly if hasattr(tensor, "__dlpack_device__") and not hasattr(tensor, "__dlpack__"): self._dlpack_data = tensor.__dlpack_device__() + elif enable_tvm_ffi: + import tvm_ffi + + self._tvm_ffi_tensor = tvm_ffi.from_dlpack(tensor) + self._dlpack_data = self._tvm_ffi_tensor.__dlpack__() else: try: # we expect no stream sync. Because torch has different default behavior @@ -149,11 +154,6 @@ class _Tensor(Tensor): self._dlpack_data = tensor.__dlpack__(stream=-1) except Exception: self._dlpack_data = tensor.__dlpack__() - if enable_tvm_ffi: - import tvm_ffi - - self._tvm_ffi_tensor = tvm_ffi.from_dlpack(tensor) - self._dlpack_data = self._tvm_ffi_tensor.__dlpack__() self._dltensor_wrapper = None self._assumed_align = assumed_align diff --git a/python/CuTeDSL/cutlass/cutlass_dsl/cutlass_ast_decorators.py b/python/CuTeDSL/cutlass/cutlass_dsl/cutlass_ast_decorators.py index 10f5c7ab..214eb710 100644 --- a/python/CuTeDSL/cutlass/cutlass_dsl/cutlass_ast_decorators.py +++ b/python/CuTeDSL/cutlass/cutlass_dsl/cutlass_ast_decorators.py @@ -647,7 +647,7 @@ def _while_execute_dynamic( def _ifexp_execute_dynamic( pred: "ir.Value", - generator_targets: tuple, + block_args: tuple, then_block: Callable, else_block: Callable, ): @@ -663,8 +663,8 @@ def _ifexp_execute_dynamic( ---------- pred : ir.Value The predicate value (a boolean IR value) that determines which branch is executed. - generator_targets : tuple - The generator targets that are passed to the then and else blocks. + block_args : tuple + The block arguments that are passed to the then and else blocks. then_block : Callable A Python function that executes the 'then' branch and returns the result(s). This will be executed if `pred` evaluates to True. @@ -698,13 +698,13 @@ def _ifexp_execute_dynamic( with ir.InsertionPoint(execution_region.region.blocks[0]): # Call the then block and unpack its results to IR values and tree structure then_results = ScfGenerator._normalize_region_result_to_list( - then_block(*generator_targets) + then_block(*block_args) ) ir_values, then_tree = cutlass_dsl.unpack_to_irvalue(then_results, "ifexp", 0) # Call the else block and unpack its results to IR values and tree structure else_results = ScfGenerator._normalize_region_result_to_list( - else_block(*generator_targets) + else_block(*block_args) ) _, else_tree = cutlass_dsl.unpack_to_irvalue(else_results, "ifexp", 0) @@ -739,11 +739,11 @@ def _ifexp_execute_dynamic( # SCF region builder for then block def then_builder(*args): # Just call the then_block as no arguments are passed to it - return then_block(*generator_targets) + return then_block(*block_args) # SCF region builder for else block def else_builder(*args): - return else_block(*generator_targets) + return else_block(*block_args) # Prepare the list of region builders for the SCF IfOp: first for "then", then for "else" region_builders = [then_builder, else_builder] diff --git a/python/CuTeDSL/cutlass/cutlass_dsl/tvm_ffi_provider.py b/python/CuTeDSL/cutlass/cutlass_dsl/tvm_ffi_provider.py index 7cec8afb..c8a53047 100644 --- a/python/CuTeDSL/cutlass/cutlass_dsl/tvm_ffi_provider.py +++ b/python/CuTeDSL/cutlass/cutlass_dsl/tvm_ffi_provider.py @@ -129,16 +129,13 @@ class TVMFFICuteCallProvider(DynamicParamPackCallProvider): cuda_global_state_ptr = self.address_of( self.cuda_global_state_symbol, self.ptr_type ) - - cuda_init_ptr = context.builder.get_or_load_global_func_ptr_from_text( - current_block, "cuda_init" - ) - cuda_load_to_device_ptr = context.builder.get_or_load_global_func_ptr_from_text( - current_block, "cuda_load_to_device" - ) - set_error_ptr = context.builder.get_or_load_global_func_ptr_from_text( - current_block, "TVMFFIErrorSetRaisedFromCStr" - ) + cuda_init_ptr = self.address_of("cuda_init", self.ptr_type) + cuda_load_to_device_ptr = self.address_of( + "cuda_load_to_device", self.ptr_type + ) + set_error_ptr = self.address_of( + "TVMFFIErrorSetRaisedFromCStr", self.ptr_type + ) with ir.InsertionPoint(current_block): # Call the callback function with the loaded ptr value diff --git a/python/CuTeDSL/cutlass/utils/dynamic_persistent_tile_scheduler.py b/python/CuTeDSL/cutlass/utils/dynamic_persistent_tile_scheduler.py index 9e6ea77f..f18a4ff3 100644 --- a/python/CuTeDSL/cutlass/utils/dynamic_persistent_tile_scheduler.py +++ b/python/CuTeDSL/cutlass/utils/dynamic_persistent_tile_scheduler.py @@ -129,7 +129,7 @@ class ClcDynamicPersistentTileScheduler: :param num_tiles_executed: Counter for executed tiles. :type num_tiles_executed: Int32 :param clc_response_ptr: Pointer of the clc rsponse. - :type clc_response_ptr: Tuple[Integer, Integer, Integer, Integer] + :type clc_response_ptr: cute.Pointer :param block_idx: The block index. :type block_idx: Tuple[Integer, Integer, Integer] """ @@ -238,7 +238,7 @@ class ClcDynamicPersistentTileScheduler: @dsl_user_op def work_tile_info_from_clc_response( - self, result_addr: Int32, *, loc=None, ip=None + self, result_addr: cute.Pointer, *, loc=None, ip=None ) -> WorkTileInfo: """ Simulates parsing CLC response data in Python. diff --git a/python/CuTeDSL/cutlass/utils/gemm/sm100.py b/python/CuTeDSL/cutlass/utils/gemm/sm100.py index 3a32646f..b6f151aa 100644 --- a/python/CuTeDSL/cutlass/utils/gemm/sm100.py +++ b/python/CuTeDSL/cutlass/utils/gemm/sm100.py @@ -293,6 +293,7 @@ def epilogue( acc_pipeline: pipeline.PipelineAsync, tCcC_base: cute.Tensor = None, mC_mnl: cute.Tensor = None, + overlapping_accum: Constexpr = False, ) -> pipeline.PipelineState: """ Epilogue function that stores accumulator results directly to global memory. @@ -310,12 +311,18 @@ def epilogue( :type epi_tile: cute.Tile :param epilogue_op: Optional elementwise operation to apply :type epilogue_op: Constexpr - :param alignment_bytes: Alignment bytes for global memory store - :type alignment_bytes: int + :param mma_tile_coord_mnl: MMA tile coordinates (M, N, L) + :type mma_tile_coord_mnl: Tuple[Int32, Int32, Int32] + :param acc_consumer_state: Accumulator consumer pipeline state + :type acc_consumer_state: pipeline.PipelineState + :param acc_pipeline: Accumulator pipeline for async operations + :type acc_pipeline: pipeline.PipelineAsync :param tCcC_base: Identity/coordinate tensor C :type tCcC_base: cute.Tensor :param mC_mnl: Global memory tensor C (full tensor for predicate computation) :type mC_mnl: cute.Tensor + :param overlapping_accum: Whether to use overlapping accumulator + :type overlapping_accum: Constexpr """ # Layout transformation for tCgC_base @@ -399,9 +406,16 @@ def epilogue( ] tTR_cC = cute.group_modes(tTR_cC, 3, cute.rank(tTR_cC)) + # Get accumulator stage index + if const_expr(overlapping_accum): + acc_stage_index = acc_consumer_state.phase + reverse_subtile = acc_stage_index == 0 + else: + acc_stage_index = acc_consumer_state.index + # Set tensor memory buffer for current tile # (T2R, T2R_M, T2R_N, EPI_M, EPI_M) - tTR_tAcc = tTR_tAcc_base[(None, None, None, None, None, acc_consumer_state.index)] + tTR_tAcc = tTR_tAcc_base[(None, None, None, None, None, acc_stage_index)] # # Wait for accumulator buffer full @@ -415,21 +429,38 @@ def epilogue( # subtile_cnt = cute.size(tTR_tAcc.shape, mode=[3]) for subtile_idx in range(subtile_cnt): + # Compute the actual subtile index + real_subtile_idx = subtile_idx + if const_expr(overlapping_accum): + if reverse_subtile: + real_subtile_idx = subtile_cnt - 1 - subtile_idx # # Get the destination and coordinate slices for this subtile # - tTR_gC_subtile = tTR_gC[(None, None, None, subtile_idx)] + tTR_gC_subtile = tTR_gC[(None, None, None, real_subtile_idx)] # # Load accumulator from tensor memory buffer to register # - tTR_tAcc_mn = tTR_tAcc[(None, None, None, subtile_idx)] + tTR_tAcc_mn = tTR_tAcc[(None, None, None, real_subtile_idx)] cute.copy(tiled_copy_t2r, tTR_tAcc_mn, tTR_rAcc) + + # # Async arrive accumulator buffer empty - # Release early for perf - if subtile_idx == subtile_cnt - 1: - with cute.arch.elect_one(): - acc_pipeline.consumer_release(acc_consumer_state) - acc_consumer_state.advance() + # + if const_expr(overlapping_accum): + # Early release when overlapping: release after processing the + # overlapping region (SF columns) so they can be reused + if subtile_idx == gemm_kernel.iter_acc_early_release_in_epilogue: + cute.arch.fence_view_async_tmem_load() + with cute.arch.elect_one(): + acc_pipeline.consumer_release(acc_consumer_state) + acc_consumer_state.advance() + else: + # Release early for perf at the last subtile + if subtile_idx == subtile_cnt - 1: + with cute.arch.elect_one(): + acc_pipeline.consumer_release(acc_consumer_state) + acc_consumer_state.advance() # # Convert to C type @@ -440,7 +471,7 @@ def epilogue( if const_expr(use_predication): # compute predicate - tTR_cC_subtile = tTR_cC[(None, None, None, subtile_idx)] + tTR_cC_subtile = tTR_cC[(None, None, None, real_subtile_idx)] pred_C_shape = (1, *tTR_cC_subtile.shape[1:]) pred_C = cute.make_rmem_tensor(pred_C_shape, Boolean) for m_idx in range(tTR_cC_subtile.shape[1]): diff --git a/python/CuTeDSL/cutlass/utils/hardware_info.py b/python/CuTeDSL/cutlass/utils/hardware_info.py index bd9f37e8..68f05235 100644 --- a/python/CuTeDSL/cutlass/utils/hardware_info.py +++ b/python/CuTeDSL/cutlass/utils/hardware_info.py @@ -41,7 +41,23 @@ class HardwareInfo: self.driver_version = self._checkCudaErrors(driver.cuDriverGetVersion()) # Getting the max active clusters for a given cluster size - def get_max_active_clusters(self, cluster_size: int) -> int: + def get_max_active_clusters( + self, cluster_size: int, stream: driver.CUstream = None + ) -> int: + """ + Get the maximum number of active clusters for a given cluster size. + + When a stream from a green context is provided, the occupancy calculation + will reflect the reduced SM partition of the green context. + + :param cluster_size: Number of blocks per cluster (must be between 1 and 32) + :type cluster_size: int + :param stream: Optional CUDA stream handle. If provided (especially from a green context), + the occupancy calculation reflects the stream's SM partition. + :type stream: driver.CUstream, optional + :return: Maximum number of active clusters + :rtype: int + """ if self._cuda_driver_version_lt(11, 8): raise RuntimeError( "CUDA Driver version < 11.8, cannot get _max_active_clusters" @@ -94,6 +110,13 @@ class HardwareInfo: launch_config.blockDimY = 1 launch_config.blockDimZ = 1 launch_config.sharedMemBytes = max_dynamic_shared_memory + + # IMPORTANT: Set the stream for green context support + # When hStream is set, cuOccupancyMaxActiveClusters will use the context + # associated with that stream, which includes the green context's SM partition + if stream is not None: + launch_config.hStream = stream + launch_config.numAttrs = 1 # max possible cluster size is 32 cluster_dims_attr = driver.CUlaunchAttribute() diff --git a/python/CuTeDSL/pyproject.toml b/python/CuTeDSL/pyproject.toml index 66af717d..5ae73380 100644 --- a/python/CuTeDSL/pyproject.toml +++ b/python/CuTeDSL/pyproject.toml @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: LicenseRef-NvidiaProprietary # # NVIDIA CORPORATION, its affiliates and licensors retain all intellectual diff --git a/python/CuTeDSL/requirements-cu13.txt b/python/CuTeDSL/requirements-cu13.txt new file mode 100644 index 00000000..b49a23e7 --- /dev/null +++ b/python/CuTeDSL/requirements-cu13.txt @@ -0,0 +1,3 @@ +# Use `pip install -r requirements-cu13.txt` with the present file to install a +# wheel consistent with the present state of the github repository +nvidia-cutlass-dsl[cu13]==4.4.0 diff --git a/python/CuTeDSL/requirements.txt b/python/CuTeDSL/requirements.txt index 29a8c9bb..bba4fe6d 100644 --- a/python/CuTeDSL/requirements.txt +++ b/python/CuTeDSL/requirements.txt @@ -1,3 +1,3 @@ # Use `pip install -r requirements.txt` with the present file to install a # wheel consistent with the present state of the github repository -nvidia-cutlass-dsl==4.4.0.dev1 +nvidia-cutlass-dsl==4.4.0 diff --git a/python/CuTeDSL/setup.sh b/python/CuTeDSL/setup.sh new file mode 100755 index 00000000..5428ece9 --- /dev/null +++ b/python/CuTeDSL/setup.sh @@ -0,0 +1,85 @@ +#!/bin/bash + +################################################################################################# +# +# Copyright (c) 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. +# +################################################################################################# + +set -e + +# Get the directory where this script is located +SCRIPT_DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" && pwd )" + +# Default to requirements.txt +REQUIREMENTS_FILE="requirements.txt" + +# Parse command line arguments +if [ $# -gt 0 ]; then + case "$1" in + --cu12) + REQUIREMENTS_FILE="requirements.txt" + echo "Installing CUDA 12 requirements..." + ;; + --cu13) + REQUIREMENTS_FILE="requirements-cu13.txt" + echo "Installing CUDA 13 requirements..." + ;; + --help|-h) + echo "Usage: $0 [--cu12|--cu13]" + echo " --cu12 Install requirements for CUDA 12 (default)" + echo " --cu13 Install requirements for CUDA 13" + exit 0 + ;; + *) + echo "Error: Unknown argument '$1'" + echo "Usage: $0 [--cu12|--cu13]" + exit 1 + ;; + esac +else + echo "Installing default requirements (CUDA 12)..." +fi + +# Check if requirements file exists +REQUIREMENTS_PATH="${SCRIPT_DIR}/${REQUIREMENTS_FILE}" +if [ ! -f "$REQUIREMENTS_PATH" ]; then + echo "Error: Requirements file not found: $REQUIREMENTS_PATH" + exit 1 +fi + +# Uninstall previous version of the CUTLASS DSL +echo "Trying to uninstall previous version of the CUTLASS DSL..." +pip uninstall nvidia-cutlass-dsl nvidia-cutlass-dsl-libs-base nvidia-cutlass-dsl-libs-cu13 -y + +# Install requirements +echo "Installing from: $REQUIREMENTS_FILE" +pip install -r "$REQUIREMENTS_PATH" + +echo "Installation complete!" diff --git a/python/cutlass_library/heuristics_provider.py b/python/cutlass_library/heuristics_provider.py index e2d60437..aefd576e 100644 --- a/python/cutlass_library/heuristics_provider.py +++ b/python/cutlass_library/heuristics_provider.py @@ -54,6 +54,7 @@ class MatmulHeuristics: def __init__(self, gpu = None): import nvMatmulHeuristics + import inspect self.mmh_lib = nvMatmulHeuristics self.gpu = gpu @@ -62,13 +63,63 @@ class MatmulHeuristics: else: nvmmhInterfaceEx = self.mmh_lib.NvMatmulHeuristicsInterfaceEx - self.lh = nvmmhInterfaceEx( + # nvidia-matmul-heuristics 0.1.0.28 changed the API: + # - Constructor: removed 'load_discovery_implicitly' and 'gpu' params + # - GPU: now set via createHardwareDescriptor() + setHardwarePredefinedGpu() + # - setBackendValueProperty renamed to setBackendPropertyValue (simpler signature) + # - getEx: added hardware_descriptor parameter + init_params = set(inspect.signature(self.mmh_lib.NvMatmulHeuristicsInterfaceEx.__init__).parameters.keys()) + self._legacy_api = 'load_discovery_implicitly' in init_params + + init_kwargs = dict( backend=self.mmh_lib.NvMatmulHeuristicsTarget["CUTLASS3"], flags=self.mmh_lib.NvMatmulHeuristicsFlags.PERF_MODEL_BASED_AUTO_TUNING, - load_discovery_implicitly=True, - gpu=self.mmh_lib.NvMatmulHeuristicsNvidiaGpu[self.gpu] if self.gpu else None ) + + if self._legacy_api: + # <= 0.1.0.27 + init_kwargs['gpu'] = self.mmh_lib.NvMatmulHeuristicsNvidiaGpu[self.gpu] if self.gpu else None + init_kwargs['load_discovery_implicitly'] = True + + self.lh = nvmmhInterfaceEx(**init_kwargs) + + # >= 0.1.0.28: gpu is set via hardware descriptor after construction, + # and passed to getEx() calls + self.hw_desc = None + if not self._legacy_api and self.gpu: + self.hw_desc = self.lh.createHardwareDescriptor() + if self.hw_desc is None: + raise RuntimeError("Failed to create hardware descriptor for GPU: " + self.gpu) + self.lh.setHardwarePredefinedGpu(self.hw_desc, self.mmh_lib.NvMatmulHeuristicsNvidiaGpu[self.gpu]) + self.backend = self.lh.createBackend(self.mmh_lib.NvMatmulHeuristicsTarget["CUTLASS3"]) + + if not self._legacy_api: + lh = self.lh + original_del = type(lh).__del__ + + def _safe_del(self_lh): + try: + original_del(self_lh) + except Exception: + pass + + type(lh).__del__ = _safe_del + + def __del__(self): + """Clean up resources in correct order before the library's __del__ runs.""" + try: + if hasattr(self, 'backend') and self.backend: + self.lh.destroyBackend(self.backend) + self.backend = None + if hasattr(self, 'hw_desc') and self.hw_desc: + self.lh.destroyHardwareDescriptor(self.hw_desc) + self.hw_desc = None + # Null out the handle so the library's __del__ skips nvMatmulHeuristicsDestroy + if hasattr(self, 'lh') and self.lh and hasattr(self.lh, 'handle'): + self.lh.handle = None + except Exception: + pass def _layout_from_cutlass(self, layouts): assert(len(layouts)==3) @@ -98,41 +149,45 @@ class MatmulHeuristics: else: return a_c + dtype_to_cublas[dtype_b] + dtype_to_cublas[dtype_c] + dtype_to_cublas[dtype_compute] + dtype_to_cublas[dtype_d] + def _set_backend_property(self, property, value): + """Compat wrapper: setBackendValueProperty (<=0.1.0.27) vs setBackendPropertyValue (>=0.1.0.28)""" + if self._legacy_api: + c_val = ctypes.c_int(value) + self.lh.setBackendValueProperty( + self.backend, property, + ctypes.byref(c_val), ctypes.sizeof(c_val) + ) + else: + self.lh.setBackendPropertyValue(self.backend, property, value) + def set_cta_div_n(self, div_n): - cta_n_div_requirement = ctypes.c_int(div_n) - self.lh.setBackendValueProperty( - self.backend, - self.mmh_lib.NvMatmulHeuristicsBackendProperty.CTA_TILE_N_DIV_REQUIREMENT, - ctypes.byref(cta_n_div_requirement), - ctypes.sizeof(cta_n_div_requirement) - ) + self._set_backend_property( + self.mmh_lib.NvMatmulHeuristicsBackendProperty.CTA_TILE_N_DIV_REQUIREMENT, div_n) def set_cta_div_m(self, div_m): - cta_m_div_requirement = ctypes.c_int(div_m) - self.lh.setBackendValueProperty( - self.backend, - self.mmh_lib.NvMatmulHeuristicsBackendProperty.CTA_TILE_M_DIV_REQUIREMENT, - ctypes.byref(cta_m_div_requirement), - ctypes.sizeof(cta_m_div_requirement) - ) + self._set_backend_property( + self.mmh_lib.NvMatmulHeuristicsBackendProperty.CTA_TILE_M_DIV_REQUIREMENT, div_m) def get_configs(self, m, n, k, batch_count, dtypes, layouts, align_a, align_b, voidC=False, use_fast_acc=True, count=1): - if use_fast_acc: - disable_fast_acc_for_fp8 = ctypes.c_int(0) - else: - disable_fast_acc_for_fp8 = ctypes.c_int(1) - self.lh.setBackendValueProperty( - self.backend, + self._set_backend_property( self.mmh_lib.NvMatmulHeuristicsBackendProperty.DISABLE_FAST_ACC_FOR_FP8, - ctypes.byref(disable_fast_acc_for_fp8), - ctypes.sizeof(disable_fast_acc_for_fp8) + 0 if use_fast_acc else 1 ) precision = self._precision_from_cutlass_dtypes(dtypes) layout = self._layout_from_cutlass(layouts) - matmul_problem = self.lh.makeNvMatmulHeuristicsProblem(m, n, k, layout, batch_count) - configs = self.lh.getEx(matmul_problem, count, self.backend, precision=precision) + if self._legacy_api: + matmul_problem = self.lh.makeNvMatmulHeuristicsProblem(m, n, k, layout, batch_count) + else: + # >= 0.1.0.28: takes (m,n,k) as a tuple + matmul_problem = self.lh.makeNvMatmulHeuristicsProblem((m, n, k), layout, batch_count) + + getEx_kwargs = dict(precision=precision) + if not self._legacy_api: + # >= 0.1.0.28: pass hardware descriptor to getEx + getEx_kwargs['hardware_descriptor'] = self.hw_desc + configs = self.lh.getEx(matmul_problem, count, self.backend, **getEx_kwargs) ret = [] for c in configs: diff --git a/test/unit/gemm/device/sm100_tensorop_gemm/CMakeLists.txt b/test/unit/gemm/device/sm100_tensorop_gemm/CMakeLists.txt index 70f6be42..4e57a807 100644 --- a/test/unit/gemm/device/sm100_tensorop_gemm/CMakeLists.txt +++ b/test/unit/gemm/device/sm100_tensorop_gemm/CMakeLists.txt @@ -70,3 +70,4 @@ cutlass_test_unit_gemm_device_add_executable( endif() add_subdirectory(narrow_precision) +add_subdirectory(extra_tests) diff --git a/test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/CMakeLists.txt b/test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/CMakeLists.txt new file mode 100644 index 00000000..38b7124e --- /dev/null +++ b/test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/CMakeLists.txt @@ -0,0 +1,95 @@ +# 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. + +if (CUTLASS_NVCC_ARCHS MATCHES 100a) +add_custom_target( + cutlass_test_unit_gemm_device_sm100_tensorop_extra + DEPENDS + cutlass_test_unit_gemm_device_tensorop_runtime_datatype_sm100 + cutlass_test_unit_gemm_device_tensorop_epilogue_fusion_sm100 + cutlass_test_unit_gemm_device_fp8_tensorop_epilogue_fusion_sm100 + cutlass_test_unit_gemm_device_sm100_dense_and_bs_gemm_stage + cutlass_test_unit_gemm_device_tensorop_stride_batch_alpha_beta_sm100 +) + +cutlass_test_unit_gemm_device_add_executable_split_file( + cutlass_test_unit_gemm_device_tensorop_runtime_datatype_sm100 + + # No batching of source to control compiler memory usage + BATCH_SOURCES ON + BATCH_SIZE 1 + + sm100_gemm_f8_f8_f8_tensor_op_f32_runtime_datatype.cu + sm100_gemm_f6_f6_f32_tensor_op_f32_runtime_datatype.cu + sm100_gemm_f4_f4_f32_tensor_op_f32_runtime_datatype.cu + sm100_gemm_f8_f4_f32_tensor_op_f32_runtime_datatype.cu +) + +cutlass_test_unit_gemm_device_add_executable_split_file( + cutlass_test_unit_gemm_device_tensorop_epilogue_fusion_sm100 + + # No batching of source to control compiler memory usage + BATCH_SOURCES ON + BATCH_SIZE 1 + + sm100_gemm_i8_i8_i8_tensor_op_s32_bias_relu.cu + sm100_gemm_i8_i8_i8_tensor_op_s32_vector_alpha_beta.cu +) + +cutlass_test_unit_gemm_device_add_executable_split_file( + cutlass_test_unit_gemm_device_fp8_tensorop_epilogue_fusion_sm100 + + # No batching of source to control compiler memory usage + BATCH_SOURCES ON + BATCH_SIZE 1 + + sm100_gemm_f8_f8_f8_tensor_op_f32_bias_relu.cu + sm100_gemm_f8_f8_f8_tensor_op_f32_bias_gelu.cu + sm100_gemm_f8_f8_f8_tensor_op_f32_bias_gelu_amax_aux.cu +) + +cutlass_test_unit_gemm_device_add_executable_split_file( + cutlass_test_unit_gemm_device_sm100_dense_and_bs_gemm_stage + + BATCH_SOURCES ON + BATCH_SIZE 1 + + sm100_gemm_f8_f8_f32_void_f8_stage.cu + sm100_gemm_f32_f32_f32_void_f32_stage.cu +) + +cutlass_test_unit_gemm_device_add_executable_split_file( + cutlass_test_unit_gemm_device_tensorop_stride_batch_alpha_beta_sm100 + + # No batching of source to control compiler memory usage + BATCH_SOURCES ON + BATCH_SIZE 1 + + sm100_gemm_f8_f8_f8_tensor_op_s32_batch_alpha_beta.cu +) +endif() diff --git a/test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_bf16_bf16_f32_tensor_op_f32.cu b/test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_bf16_bf16_f32_tensor_op_f32.cu new file mode 100644 index 00000000..c2827b64 --- /dev/null +++ b/test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_bf16_bf16_f32_tensor_op_f32.cu @@ -0,0 +1,310 @@ +/*************************************************************************************************** + * 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 Tests for device-wide GEMM interface +*/ + +#include + +#include "cutlass/cutlass.h" +#include "cute/tensor.hpp" +#include "cute/atom/mma_atom.hpp" + +#include "cutlass/numeric_types.h" + +#include "cutlass/gemm/device/gemm_universal_adapter.h" +#include "cutlass/gemm/kernel/gemm_universal.hpp" +#include "cutlass/gemm/collective/collective_builder.hpp" +#include "cutlass/gemm/dispatch_policy.hpp" +#include "cutlass/epilogue/dispatch_policy.hpp" +#include "cutlass/epilogue/collective/collective_builder.hpp" +#include "cutlass/epilogue/thread/activation.h" + +#include "../../../../common/cutlass_unit_test.h" +#include "../../gemm_testbed_3x.hpp" + +using namespace cute; + +#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) + +/// A Row B Col +TEST(SM100_Device_Gemm_f16t_f16n_f32t_tensorop_2sm_f32, 512x512x128_4x4x1) { + using ElementA = cutlass::bfloat16_t; + using ElementB = cutlass::bfloat16_t; + using ElementC = void; + using ElementD = float; + using ElementCompute = float; + using ElementAccumulator = float; + using GmemLayoutA = cutlass::layout::RowMajor; + using GmemLayoutB = cutlass::layout::ColumnMajor; + using GmemLayoutC = cutlass::layout::RowMajor; + using MmaTileShape_MNK = Shape<_256,_128,_128>; + using ClusterShape_MNK = Shape<_4,_4,_1>; + + // + // Construct CollectiveEpilogue + // + + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + MmaTileShape_MNK, ClusterShape_MNK, + cutlass::epilogue::collective::EpilogueTileAuto, + ElementAccumulator, ElementCompute, + ElementC, GmemLayoutC, 16, + ElementD, GmemLayoutC, 16, + cutlass::epilogue::TmaWarpSpecialized2Sm + >::CollectiveOp; + + // + // Construct CollectiveMainloop + // + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + ElementA, GmemLayoutA, 8, + ElementB, GmemLayoutB, 8, + ElementAccumulator, + MmaTileShape_MNK, ClusterShape_MNK, + cutlass::gemm::collective::StageCountAutoCarveout(sizeof(typename CollectiveEpilogue::SharedStorage))>, + cutlass::gemm::KernelTmaWarpSpecialized2SmSm100 + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, + CollectiveMainloop, + CollectiveEpilogue + >; + + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + auto pass = test::gemm::device::TestSmallFusion(1.0, 0); + EXPECT_TRUE(pass); +} + +/// A Col B Row +TEST(SM100_Device_Gemm_f16n_f16t_f32t_tensorop_2sm_f32, 512x512x128_4x4x1) { + using ElementA = cutlass::bfloat16_t; + using ElementB = cutlass::bfloat16_t; + using ElementC = void; + using ElementD = float; + using ElementCompute = float; + using ElementAccumulator = float; + using GmemLayoutA = cutlass::layout::ColumnMajor; + using GmemLayoutB = cutlass::layout::RowMajor; + using GmemLayoutC = cutlass::layout::RowMajor; + using MmaTileShape_MNK = Shape<_256,_128,_128>; + using ClusterShape_MNK = Shape<_4,_4,_1>; + + // + // Construct CollectiveEpilogue + // + + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + MmaTileShape_MNK, ClusterShape_MNK, + cutlass::epilogue::collective::EpilogueTileAuto, + ElementAccumulator, ElementCompute, + ElementC, GmemLayoutC, 16, + ElementD, GmemLayoutC, 16, + cutlass::epilogue::TmaWarpSpecialized2Sm + >::CollectiveOp; + + // + // Construct CollectiveMainloop + // + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + ElementA, GmemLayoutA, 8, + ElementB, GmemLayoutB, 8, + ElementAccumulator, + MmaTileShape_MNK, ClusterShape_MNK, + cutlass::gemm::collective::StageCountAutoCarveout(sizeof(typename CollectiveEpilogue::SharedStorage))>, + cutlass::gemm::KernelTmaWarpSpecialized2SmSm100 + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, + CollectiveMainloop, + CollectiveEpilogue + >; + + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + auto pass = test::gemm::device::TestSmallFusion(1.0, 0); + EXPECT_TRUE(pass); +} + +/// A Row B Row +TEST(SM100_Device_Gemm_f16t_f16t_f32t_tensorop_2sm_f32, 512x512x128_4x4x1) { + using ElementA = cutlass::bfloat16_t; + using ElementB = cutlass::bfloat16_t; + using ElementC = void; + using ElementD = float; + using ElementCompute = float; + using ElementAccumulator = float; + using GmemLayoutA = cutlass::layout::RowMajor; + using GmemLayoutB = cutlass::layout::RowMajor; + using GmemLayoutC = cutlass::layout::RowMajor; + using MmaTileShape_MNK = Shape<_256,_128,_128>; + using ClusterShape_MNK = Shape<_4,_4,_1>; + + // + // Construct CollectiveEpilogue + // + + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + MmaTileShape_MNK, ClusterShape_MNK, + cutlass::epilogue::collective::EpilogueTileAuto, + ElementAccumulator, ElementCompute, + ElementC, GmemLayoutC, 16, + ElementD, GmemLayoutC, 16, + cutlass::epilogue::TmaWarpSpecialized2Sm + >::CollectiveOp; + + // + // Construct CollectiveMainloop + // + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + ElementA, GmemLayoutA, 8, + ElementB, GmemLayoutB, 8, + ElementAccumulator, + MmaTileShape_MNK, ClusterShape_MNK, + cutlass::gemm::collective::StageCountAutoCarveout(sizeof(typename CollectiveEpilogue::SharedStorage))>, + cutlass::gemm::KernelTmaWarpSpecialized2SmSm100 + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, + CollectiveMainloop, + CollectiveEpilogue + >; + + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + auto pass = test::gemm::device::TestSmallFusion(1.0, 0); + EXPECT_TRUE(pass); +} + +/// A Col B Col +TEST(SM100_Device_Gemm_f16n_f16n_f32t_tensorop_2sm_f32, 512x512x128_4x4x1) { + using ElementA = cutlass::bfloat16_t; + using ElementB = cutlass::bfloat16_t; + using ElementC = void; + using ElementD = float; + using ElementCompute = float; + using ElementAccumulator = float; + using GmemLayoutA = cutlass::layout::ColumnMajor; + using GmemLayoutB = cutlass::layout::ColumnMajor; + using GmemLayoutC = cutlass::layout::RowMajor; + using MmaTileShape_MNK = Shape<_256,_128,_128>; + using ClusterShape_MNK = Shape<_4,_4,_1>; + + // + // Construct CollectiveEpilogue + // + + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + MmaTileShape_MNK, ClusterShape_MNK, + cutlass::epilogue::collective::EpilogueTileAuto, + ElementAccumulator, ElementCompute, + ElementC, GmemLayoutC, 16, + ElementD, GmemLayoutC, 16, + cutlass::epilogue::TmaWarpSpecialized2Sm + >::CollectiveOp; + + // + // Construct CollectiveMainloop + // + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + ElementA, GmemLayoutA, 8, + ElementB, GmemLayoutB, 8, + ElementAccumulator, + MmaTileShape_MNK, ClusterShape_MNK, + cutlass::gemm::collective::StageCountAutoCarveout(sizeof(typename CollectiveEpilogue::SharedStorage))>, + cutlass::gemm::KernelTmaWarpSpecialized2SmSm100 + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, + CollectiveMainloop, + CollectiveEpilogue + >; + + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + auto pass = test::gemm::device::TestSmallFusion(1.0, 0); + EXPECT_TRUE(pass); +} + +TEST(SM100_Device_Gemm_bf16t_bf16t_bf32_void_f32n_tensor_op, 128x256x64_1x2x1) { + using ElementA = cutlass::bfloat16_t; + using LayoutA = cutlass::layout::RowMajor; + using ElementB = cutlass::bfloat16_t; + using LayoutB = cutlass::layout::RowMajor; + using ElementAccumulator = float; + using LayoutC = cutlass::layout::ColumnMajor; + using MmaTileShape = Shape<_128,_128,_64>; + using ClusterShape = Shape<_1,_2,_1>; + + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + MmaTileShape, ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + float, float, + void, LayoutC, 8, + float, LayoutC, 8, + cutlass::epilogue::TmaWarpSpecialized1Sm + >::CollectiveOp; + + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + cutlass::half_t, LayoutA, 8, + cutlass::half_t, LayoutB, 8, + float, + MmaTileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout< + static_cast(sizeof(typename CollectiveEpilogue::SharedStorage))>, + cutlass::gemm::KernelTmaWarpSpecialized1SmSm100 + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, + CollectiveMainloop, + CollectiveEpilogue + >; + + using namespace test::gemm::device; + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + auto pass = test::gemm::device::TestSmall(1.0, 0.0); + EXPECT_TRUE(pass); +} + +#endif // #if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) diff --git a/test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_f16_f16_f16_tensor_op_f32_b2b.cu b/test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_f16_f16_f16_tensor_op_f32_b2b.cu new file mode 100644 index 00000000..1cd1ab3a --- /dev/null +++ b/test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_f16_f16_f16_tensor_op_f32_b2b.cu @@ -0,0 +1,257 @@ +/*************************************************************************************************** + * 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 Tests for device-wide SM100 back-to-back GEMM interface +*/ + + +#include + +#include "cutlass/cutlass.h" +#include "cute/tensor.hpp" +#include "cute/atom/mma_atom.hpp" + +#include "cutlass/numeric_types.h" + +#include "cutlass/gemm/device/gemm_universal_adapter.h" +#include "cutlass/gemm/collective/collective_builder.hpp" +#include "cutlass/epilogue/collective/collective_builder.hpp" +#include "cutlass/epilogue/collective/sm70_epilogue_vectorized.hpp" +#include "cutlass/epilogue/collective/default_epilogue.hpp" +#include "cutlass/epilogue/thread/linear_combination.h" +#include "cutlass/epilogue/thread/activation.h" + +#include "../../../../common/cutlass_unit_test.h" +#include "../../gemm_testbed_3x_b2b.hpp" + +#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) + +TEST(SM100_Device_Gemm_f16t_f16n_f16t_tensor_op_gmma_f32_b2b, 128x64x64_1x1x1_128x256x64_1x2x1) { + + using CollectiveEpilogue0 = + typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + cute::Shape, + cute::Shape, + cutlass::epilogue::collective::EpilogueTileAuto, + float, float, + cutlass::half_t, cutlass::layout::RowMajor, 8, + cutlass::half_t, cutlass::layout::RowMajor, 8, + cutlass::epilogue::TmaWarpSpecialized1Sm + >::CollectiveOp; + + using CollectiveMainloop0 = + typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + cutlass::half_t, cutlass::layout::RowMajor, 8, + cutlass::half_t, cutlass::layout::ColumnMajor, 8, + float, + cute::Shape, + cute::Shape, + cutlass::gemm::collective::StageCountAutoCarveout, + cutlass::gemm::KernelTmaWarpSpecialized1SmSm100 + >::CollectiveOp; + + using GemmKernel0 = cutlass::gemm::kernel::GemmUniversal< + cute::Shape, + CollectiveMainloop0, + CollectiveEpilogue0>; + + using CollectiveEpilogue1 = + typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + cute::Shape, + cute::Shape, + cutlass::epilogue::collective::EpilogueTileAuto, + float, float, + cutlass::half_t, cutlass::layout::RowMajor, 8, + cutlass::half_t, cutlass::layout::RowMajor, 8, + cutlass::epilogue::TmaWarpSpecialized1Sm + >::CollectiveOp; + + using CollectiveMainloop1 = + typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + cutlass::half_t, cutlass::layout::RowMajor, 8, + cutlass::half_t, cutlass::layout::ColumnMajor, 8, + float, + cute::Shape, + cute::Shape, + cutlass::gemm::collective::StageCountAutoCarveout, + cutlass::gemm::KernelTmaWarpSpecialized1SmSm100 + >::CollectiveOp; + + using GemmKernel1 = cutlass::gemm::kernel::GemmUniversal< + cute::Shape, + CollectiveMainloop1, + CollectiveEpilogue1>; + + using Gemm0 = cutlass::gemm::device::GemmUniversalAdapter; + using Gemm1 = cutlass::gemm::device::GemmUniversalAdapter; + + EXPECT_TRUE((test::gemm::device::TestAllB2B())); +} + +TEST(SM100_Device_Gemm_f16t_f16n_f16t_tensor_op_gmma_f32_b2b, 128x64x64_1x2x1_128x256x64_2x2x1) { + + using CollectiveEpilogue0 = + typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + cute::Shape, + cute::Shape, + cutlass::epilogue::collective::EpilogueTileAuto, + float, float, + cutlass::half_t, cutlass::layout::RowMajor, 8, + cutlass::half_t, cutlass::layout::RowMajor, 8, + cutlass::epilogue::TmaWarpSpecialized1Sm + >::CollectiveOp; + + using CollectiveMainloop0 = + typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + cutlass::half_t, cutlass::layout::RowMajor, 8, + cutlass::half_t, cutlass::layout::ColumnMajor, 8, + float, + cute::Shape, + cute::Shape, + cutlass::gemm::collective::StageCountAutoCarveout, + cutlass::gemm::KernelTmaWarpSpecialized1SmSm100 + >::CollectiveOp; + + using GemmKernel0 = cutlass::gemm::kernel::GemmUniversal< + cute::Shape, + CollectiveMainloop0, + CollectiveEpilogue0>; + + using CollectiveEpilogue1 = + typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + cute::Shape, + cute::Shape, + cutlass::epilogue::collective::EpilogueTileAuto, + float, float, + cutlass::half_t, cutlass::layout::RowMajor, 8, + cutlass::half_t, cutlass::layout::RowMajor, 8, + cutlass::epilogue::TmaWarpSpecialized1Sm + >::CollectiveOp; + + using CollectiveMainloop1 = + typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + cutlass::half_t, cutlass::layout::RowMajor, 8, + cutlass::half_t, cutlass::layout::ColumnMajor, 8, + float, + cute::Shape, + cute::Shape, + cutlass::gemm::collective::StageCountAutoCarveout, + cutlass::gemm::KernelTmaWarpSpecialized1SmSm100 + >::CollectiveOp; + + using GemmKernel1 = cutlass::gemm::kernel::GemmUniversal< + cute::Shape, + CollectiveMainloop1, + CollectiveEpilogue1>; + + using Gemm0 = cutlass::gemm::device::GemmUniversalAdapter; + using Gemm1 = cutlass::gemm::device::GemmUniversalAdapter; + + EXPECT_TRUE((test::gemm::device::TestAllB2B())); +} + +TEST(SM100_Device_Gemm_f16t_f16n_f16t_tensor_op_gmma_f32_b2b, 128x128x64_2x2x1_256x128x64_2x1x1) { + + using CollectiveEpilogue0 = + typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + cute::Shape, + cute::Shape, + cutlass::epilogue::collective::EpilogueTileAuto, + float, float, + cutlass::half_t, cutlass::layout::RowMajor, 8, + cutlass::half_t, cutlass::layout::RowMajor, 8, + cutlass::epilogue::TmaWarpSpecialized1Sm + >::CollectiveOp; + + using CollectiveMainloop0 = + typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + cutlass::half_t, cutlass::layout::RowMajor, 8, + cutlass::half_t, cutlass::layout::ColumnMajor, 8, + float, + cute::Shape, + cute::Shape, + cutlass::gemm::collective::StageCountAutoCarveout, + cutlass::gemm::KernelTmaWarpSpecialized1SmSm100 + >::CollectiveOp; + + using GemmKernel0 = cutlass::gemm::kernel::GemmUniversal< + cute::Shape, + CollectiveMainloop0, + CollectiveEpilogue0>; + + using CollectiveEpilogue1 = + typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + cute::Shape, + cute::Shape, + cutlass::epilogue::collective::EpilogueTileAuto, + float, float, + cutlass::half_t, cutlass::layout::RowMajor, 8, + cutlass::half_t, cutlass::layout::RowMajor, 8, + cutlass::epilogue::TmaWarpSpecialized1Sm + >::CollectiveOp; + + using CollectiveMainloop1 = + typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + cutlass::half_t, cutlass::layout::RowMajor, 8, + cutlass::half_t, cutlass::layout::ColumnMajor, 8, + float, + cute::Shape, + cute::Shape, + cutlass::gemm::collective::StageCountAutoCarveout, + cutlass::gemm::KernelTmaWarpSpecialized1SmSm100 + >::CollectiveOp; + + using GemmKernel1 = cutlass::gemm::kernel::GemmUniversal< + cute::Shape, + CollectiveMainloop1, + CollectiveEpilogue1>; + + using Gemm0 = cutlass::gemm::device::GemmUniversalAdapter; + using Gemm1 = cutlass::gemm::device::GemmUniversalAdapter; + + EXPECT_TRUE((test::gemm::device::TestAllB2B())); +} + +#endif // defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) + diff --git a/test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_f16_f16_f16_tensor_op_f32_stream_k.cu b/test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_f16_f16_f16_tensor_op_f32_stream_k.cu new file mode 100644 index 00000000..61ee4a83 --- /dev/null +++ b/test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_f16_f16_f16_tensor_op_f32_stream_k.cu @@ -0,0 +1,235 @@ +/*************************************************************************************************** + * 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. + * + **************************************************************************************************/ +/*! \file + \brief Tests for device-wide GEMM interface with stream-K scheduling +*/ + +#include + +#include "cutlass/cutlass.h" +#include "cute/tensor.hpp" +#include "cute/atom/mma_atom.hpp" + +#include "cutlass/numeric_types.h" + +#include "cutlass/gemm/device/gemm_universal_adapter.h" +#include "cutlass/gemm/kernel/gemm_universal.hpp" +#include "cutlass/gemm/kernel/tile_scheduler.hpp" +#include "cutlass/gemm/collective/collective_builder.hpp" +#include "cutlass/epilogue/collective/collective_builder.hpp" +#include "cutlass/epilogue/collective/sm70_epilogue_vectorized.hpp" +#include "cutlass/epilogue/collective/default_epilogue.hpp" +#include "cutlass/epilogue/thread/linear_combination.h" + +#include "../../../../common/cutlass_unit_test.h" +#include "../../gemm_testbed_3x.hpp" + +#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) + +using namespace cute; + +TEST(SM100_Device_Gemm_f16t_f16t_f32n_tensor_op_gmma_f32_stream_k, 128x256x64_1x2x1) { + using ElementA = cutlass::half_t; + using LayoutA = cutlass::layout::RowMajor; + using ElementB = cutlass::half_t; + using LayoutB = cutlass::layout::RowMajor; + using ElementAccumulator = float; + using LayoutC = cutlass::layout::ColumnMajor; + using MmaTileShape_MNK = Shape<_128,_128,_64>; + using ClusterShape_MNK = Shape<_1,_2,_1>; + + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + MmaTileShape_MNK, ClusterShape_MNK, + cutlass::epilogue::collective::EpilogueTileAuto, + float, float, + cutlass::half_t, LayoutC, 8, + cutlass::half_t, LayoutC, 8, + cutlass::epilogue::TmaWarpSpecialized1Sm + >::CollectiveOp; + + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + cutlass::half_t, LayoutA, 8, + cutlass::half_t, LayoutB, 8, + float, + MmaTileShape_MNK, ClusterShape_MNK, + cutlass::gemm::collective::StageCountAutoCarveout< + static_cast(sizeof(typename CollectiveEpilogue::SharedStorage))>, + cutlass::gemm::KernelTmaWarpSpecialized1SmSm100 + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, + CollectiveMainloop, + CollectiveEpilogue, + cutlass::gemm::StreamKScheduler + >; + + using namespace test::gemm::device; + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + using Testbed = Testbed3x; + bool result = TestSmall(1.0, 0.0, CheckEquality::EXACT, ScalarLoc::ON_DEVICE, VectorScale::ENABLED, {64, 1024, 2048}); + EXPECT_TRUE(result); +} + +TEST(SM100_Device_Gemm_f16t_f16t_f32n_tensor_op_gmma_f32_stream_k, 256x128x64_2x1x1) { + using ElementA = cutlass::half_t; + using LayoutA = cutlass::layout::RowMajor; + using ElementB = cutlass::half_t; + using LayoutB = cutlass::layout::RowMajor; + using ElementAccumulator = float; + using LayoutC = cutlass::layout::ColumnMajor; + using MmaTileShape_MNK = Shape<_128,_128,_64>; + using ClusterShape_MNK = Shape<_2,_1,_1>; + + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + MmaTileShape_MNK, ClusterShape_MNK, + cutlass::epilogue::collective::EpilogueTileAuto, + float, float, + cutlass::half_t, LayoutC, 8, + cutlass::half_t, LayoutC, 8, + cutlass::epilogue::TmaWarpSpecialized1Sm + >::CollectiveOp; + + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + cutlass::half_t, LayoutA, 8, + cutlass::half_t, LayoutB, 8, + float, + MmaTileShape_MNK, ClusterShape_MNK, + cutlass::gemm::collective::StageCountAutoCarveout< + static_cast(sizeof(typename CollectiveEpilogue::SharedStorage))>, + cutlass::gemm::KernelTmaWarpSpecialized1SmSm100 + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, + CollectiveMainloop, + CollectiveEpilogue, + cutlass::gemm::StreamKScheduler + >; + + using namespace test::gemm::device; + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + using Testbed = Testbed3x; + bool result = TestSmall(1.0, 0.0, CheckEquality::EXACT, ScalarLoc::ON_DEVICE, VectorScale::ENABLED, {64, 1024, 2048}); + EXPECT_TRUE(result); +} + +TEST(SM100_Device_Gemm_f16t_f16t_f32n_tensor_op_gmma_f32_stream_k, 256x256x64_2x2x1) { + using LayoutA = cutlass::layout::RowMajor; + using LayoutB = cutlass::layout::RowMajor; + using LayoutC = cutlass::layout::ColumnMajor; + using MmaTileShape_MNK = Shape<_128,_128,_64>; + using ClusterShape_MNK = Shape<_2,_2,_1>; + + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + MmaTileShape_MNK, ClusterShape_MNK, + cutlass::epilogue::collective::EpilogueTileAuto, + float, float, + cutlass::half_t, LayoutC, 8, + cutlass::half_t, LayoutC, 8, + cutlass::epilogue::TmaWarpSpecialized1Sm + >::CollectiveOp; + + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + cutlass::half_t, LayoutA, 8, + cutlass::half_t, LayoutB, 8, + float, + MmaTileShape_MNK, ClusterShape_MNK, + cutlass::gemm::collective::StageCountAutoCarveout< + static_cast(sizeof(typename CollectiveEpilogue::SharedStorage))>, + cutlass::gemm::KernelTmaWarpSpecialized1SmSm100 + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, + CollectiveMainloop, + CollectiveEpilogue, + cutlass::gemm::StreamKScheduler + >; + + using namespace test::gemm::device; + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + using Testbed = Testbed3x; + bool result = TestSmall(1.0, 0.0, CheckEquality::EXACT, ScalarLoc::ON_DEVICE, VectorScale::ENABLED, {64, 1024, 2048}); + EXPECT_TRUE(result); +} + +/////////////////////////////////////////////////////////////////////////////// + +TEST(SM100_Device_Gemm_f16t_f16n_f32n_tensor_op_gmma_f32_stream_k, 256x128x64_2x4x1) { + using LayoutA = cutlass::layout::RowMajor; + using LayoutB = cutlass::layout::ColumnMajor; + using LayoutC = cutlass::layout::ColumnMajor; + using MmaTileShape_MNK = Shape<_128,_64,_64>; + using ClusterShape_MNK = Shape<_2,_4,_1>; + + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + MmaTileShape_MNK, ClusterShape_MNK, + cutlass::epilogue::collective::EpilogueTileAuto, + float, float, + cutlass::half_t, LayoutC, 8, + cutlass::half_t, LayoutC, 8, + cutlass::epilogue::TmaWarpSpecialized1Sm + >::CollectiveOp; + + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + cutlass::half_t, LayoutA, 8, + cutlass::half_t, LayoutB, 8, + float, + MmaTileShape_MNK, ClusterShape_MNK, + cutlass::gemm::collective::StageCountAutoCarveout< + static_cast(sizeof(typename CollectiveEpilogue::SharedStorage))>, + cutlass::gemm::KernelTmaWarpSpecialized1SmSm100 + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, + CollectiveMainloop, + CollectiveEpilogue, + cutlass::gemm::StreamKScheduler + >; + + using namespace test::gemm::device; + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + using Testbed = Testbed3x; + bool result = TestSmall(1.0, 0.0, CheckEquality::EXACT, ScalarLoc::ON_DEVICE, VectorScale::ENABLED, {64, 1024, 2048}); + EXPECT_TRUE(result); +} + +#endif // defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) diff --git a/test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_f16_f16_f16_tensor_op_f32_swap_ab_bias_relu.cu b/test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_f16_f16_f16_tensor_op_f32_swap_ab_bias_relu.cu new file mode 100644 index 00000000..872aafea --- /dev/null +++ b/test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_f16_f16_f16_tensor_op_f32_swap_ab_bias_relu.cu @@ -0,0 +1,168 @@ +/*************************************************************************************************** + * 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 Tests for device-wide GEMM interface +*/ + +#include + +#include "cutlass/cutlass.h" +#include "cute/tensor.hpp" +#include "cute/atom/mma_atom.hpp" + +#include "cutlass/numeric_types.h" +#include "cutlass/arch/mma_sm100.h" + +#include "cutlass/gemm/device/gemm_universal_adapter.h" +#include "cutlass/gemm/kernel/gemm_universal.hpp" +#include "cutlass/gemm/collective/collective_builder.hpp" +#include "cutlass/epilogue/dispatch_policy.hpp" +#include "cutlass/epilogue/collective/collective_builder.hpp" + +#include "../../../../common/cutlass_unit_test.h" +#include "../../gemm_testbed_3x.hpp" + +using namespace cute; + +#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) + +/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// +///////////////////////////////////////////// 128x64x64 1x1x1 TMEM 4x1 //////////////////////////////////////////// +/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// + +TEST(SM100_Device_Gemm_f16t_f16n_f16t_tensorop_1cta_f32_bias_relu_aux, 128x256x64_1x1x1) { + using LayoutA = cutlass::layout::RowMajor; + using LayoutB = cutlass::layout::ColumnMajor; + using LayoutC = cutlass::layout::RowMajor; + using ElementA = cutlass::half_t; + using ElementB = cutlass::half_t; + using ElementC = cutlass::half_t; + using ElementD = cutlass::half_t; + using ElementAccumulator = float; + using ElementCompute = float; + using ElementBias = cutlass::half_t; + using ElementAux = cutlass::uint1b_t; + using MmaTileShape = cute::Shape<_128,_256,Int<128 / sizeof(ElementA)>>; + using ClusterShape = Shape<_1,_1,_1>; + + using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecialized1Sm; + using FusionOperation = cutlass::epilogue::fusion::LinCombPerColBiasEltActAux< + LayoutC, cutlass::epilogue::thread::Clamp, ElementD, ElementCompute, ElementAux, ElementBias>; + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + MmaTileShape, ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + ElementAccumulator, ElementCompute, + ElementC, LayoutC, 16 / sizeof(ElementC), + ElementD, LayoutC, 16 / sizeof(ElementD), + EpilogueSchedule, + FusionOperation + >::CollectiveOp; + + using MainloopSchedule = cutlass::gemm::KernelTmaWarpSpecialized1SmSm100; + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + ElementA, LayoutA, 16 / sizeof(ElementA), + ElementB, LayoutB, 16 / sizeof(ElementB), + ElementAccumulator, + MmaTileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout(sizeof(typename CollectiveEpilogue::SharedStorage))>, + MainloopSchedule + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, + CollectiveMainloop, + CollectiveEpilogue + >; + + using namespace test::gemm::device; + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + + auto pass = TestSmallFusion(2, 0.5, CheckEquality::EXACT); + EXPECT_TRUE(pass); +} + +TEST(SM100_Device_Gemm_f16t_f16n_f16t_tensorop_2cta_f32_bias_gelu_aux, 128x1024x64_2x4x1) { + using LayoutA = cutlass::layout::RowMajor; + using LayoutB = cutlass::layout::ColumnMajor; + using LayoutC = cutlass::layout::ColumnMajor; + using ElementA = cutlass::half_t; + using ElementB = cutlass::half_t; + using ElementC = cutlass::half_t; + using ElementD = cutlass::half_t; + using ElementAccumulator = float; + using ElementCompute = float; + using ElementAux = cutlass::half_t; + using ElementBias = cutlass::half_t; + using MmaTileShape = Shape<_128,_256,Int<128 / sizeof(ElementA)>>; + using ClusterShape = Shape<_2,_4,_1>; + + using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecialized2Sm; + using FusionOperation = cutlass::epilogue::fusion::LinCombPerColBiasEltActAux< + LayoutC, cutlass::epilogue::thread::GELU, ElementD, ElementCompute, ElementAux, ElementBias>; + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + MmaTileShape, ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + ElementAccumulator, ElementCompute, + ElementC, LayoutC, 16 / sizeof(ElementC), + ElementD, LayoutC, 16 / sizeof(ElementD), + EpilogueSchedule, + FusionOperation + >::CollectiveOp; + + using MainloopSchedule = cutlass::gemm::KernelTmaWarpSpecialized2SmSm100; + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + ElementA, LayoutA, 16 / sizeof(ElementA), + ElementB, LayoutB, 16 / sizeof(ElementB), + ElementAccumulator, + MmaTileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout(sizeof(typename CollectiveEpilogue::SharedStorage))>, + MainloopSchedule + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, + CollectiveMainloop, + CollectiveEpilogue + >; + + using namespace test::gemm::device; + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + + auto pass = TestSmallFusion(2, 0.5, CheckEquality::RELATIVE); + EXPECT_TRUE(pass); +} + +#endif // #if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) diff --git a/test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_f16_f16_f32_tensor_op_f32.cu b/test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_f16_f16_f32_tensor_op_f32.cu new file mode 100644 index 00000000..d514ba32 --- /dev/null +++ b/test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_f16_f16_f32_tensor_op_f32.cu @@ -0,0 +1,96 @@ +/*************************************************************************************************** + * 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. + * + **************************************************************************************************/ + +#include + +#include "cutlass/cutlass.h" +#include "cute/tensor.hpp" +#include "cute/atom/mma_atom.hpp" + +#include "cutlass/numeric_types.h" +#include "cutlass/gemm/device/gemm_universal_adapter.h" +#include "cutlass/gemm/kernel/gemm_universal.hpp" +#include "cutlass/gemm/collective/collective_builder.hpp" +#include "cutlass/epilogue/dispatch_policy.hpp" +#include "cutlass/epilogue/collective/collective_builder.hpp" +#include "cutlass/epilogue/thread/activation.h" + +#include "../../../../common/cutlass_unit_test.h" +#include "../../gemm_testbed_3x.hpp" + +using namespace cute; + +#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) + +TEST(SM100_Device_Gemm_f16t_f16t_f32_void_f16n_tensor_op, 128x256x64_1x2x1) { + using ElementA = cutlass::half_t; + using LayoutA = cutlass::layout::RowMajor; + using ElementB = cutlass::half_t; + using LayoutB = cutlass::layout::RowMajor; + using ElementAccumulator = float; + using LayoutC = cutlass::layout::ColumnMajor; + using MmaTileShape_MNK = Shape<_128,_128,_64>; + using ClusterShape_MNK = Shape<_1,_2,_1>; + + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + MmaTileShape_MNK, ClusterShape_MNK, + cutlass::epilogue::collective::EpilogueTileAuto, + float, float, + void, LayoutC, 8, + cutlass::half_t, LayoutC, 8, + cutlass::epilogue::TmaWarpSpecialized1Sm + >::CollectiveOp; + + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + cutlass::half_t, LayoutA, 8, + cutlass::half_t, LayoutB, 8, + float, + MmaTileShape_MNK, ClusterShape_MNK, + cutlass::gemm::collective::StageCountAutoCarveout< + static_cast(sizeof(typename CollectiveEpilogue::SharedStorage))>, + cutlass::gemm::KernelTmaWarpSpecialized1SmSm100 + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, + CollectiveMainloop, + CollectiveEpilogue + >; + + using namespace test::gemm::device; + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + auto pass = test::gemm::device::TestSmall(1.0, 0.0); + EXPECT_TRUE(pass); +} + +#endif // #if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) diff --git a/test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_f16_f16_f32_void_f16_stage.cu b/test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_f16_f16_f32_void_f16_stage.cu new file mode 100644 index 00000000..651a25c9 --- /dev/null +++ b/test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_f16_f16_f32_void_f16_stage.cu @@ -0,0 +1,233 @@ +/*************************************************************************************************** + * 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 Test for sm100 dense gemm stage +*/ +#include "cute/atom/mma_atom.hpp" +#include "cute/tensor.hpp" +#include "cutlass/arch/mma_sm100.h" +#include "cutlass/cutlass.h" +#include "cutlass/epilogue/collective/collective_builder.hpp" +#include "cutlass/epilogue/dispatch_policy.hpp" +#include "cutlass/epilogue/thread/linear_combination.h" +#include "cutlass/gemm/collective/collective_builder.hpp" +#include "cutlass/gemm/device/gemm_universal_adapter.h" +#include "cutlass/gemm/kernel/gemm_universal.hpp" +#include "cutlass/numeric_types.h" +#include "../../../../common/cutlass_unit_test.h" +#include "../../gemm_testbed_3x.hpp" + +using namespace cute; + +namespace cutlass3x_sm100_tensorop_gemm_f16_f16_f32_void_f16_128x128x64_1x1x1_0_tnt_align8_1sm_stage6 { + using CollectiveEpilogue = + typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + cute::Shape, + cute::Shape, + cutlass::epilogue::collective::EpilogueTileAuto, + float, float, + void, cutlass::layout::RowMajor, 8, + cutlass::half_t, cutlass::layout::RowMajor, 8, + cutlass::epilogue::TmaWarpSpecialized1Sm, + cutlass::epilogue::fusion::LinearCombination< + cutlass::half_t, + float, + void, + float + > + >::CollectiveOp; + + using CollectiveMainloop = + typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + cutlass::half_t, cutlass::layout::RowMajor, 8, + cutlass::half_t, cutlass::layout::ColumnMajor, 8, + float, + cute::Shape, + cute::Shape, + cutlass::gemm::collective::StageCountAutoCarveout(sizeof(typename CollectiveEpilogue::SharedStorage))>, + cutlass::gemm::KernelTmaWarpSpecialized1SmSm100 + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + cute::Shape, + CollectiveMainloop, + CollectiveEpilogue, + void>; + + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; +} + +namespace cutlass3x_sm100_tensorop_gemm_f16_f16_f32_void_f16_128x256x64_1x1x1_0_tnt_align8_1sm_stage4 { + using CollectiveEpilogue = + typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + cute::Shape, + cute::Shape, + cutlass::epilogue::collective::EpilogueTileAuto, + float, float, + void, cutlass::layout::RowMajor, 8, + cutlass::half_t, cutlass::layout::RowMajor, 8, + cutlass::epilogue::TmaWarpSpecialized1Sm, + cutlass::epilogue::fusion::LinearCombination< + cutlass::half_t, + float, + void, + float + > + >::CollectiveOp; + + using CollectiveMainloop = + typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + cutlass::half_t, cutlass::layout::RowMajor, 8, + cutlass::half_t, cutlass::layout::ColumnMajor, 8, + float, + cute::Shape, + cute::Shape, + cutlass::gemm::collective::StageCountAutoCarveout(sizeof(typename CollectiveEpilogue::SharedStorage))>, + cutlass::gemm::KernelTmaWarpSpecialized1SmSm100 + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + cute::Shape, + CollectiveMainloop, + CollectiveEpilogue, + void>; + + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; +} + +namespace cutlass3x_sm100_tensorop_gemm_f16_f16_f32_void_f16_256x128x64_2x1x1_0_tnt_align8_2sm_stage8 { + using CollectiveEpilogue = + typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + cute::Shape, + cute::Shape, + cutlass::epilogue::collective::EpilogueTileAuto, + float, float, + void, cutlass::layout::RowMajor, 8, + cutlass::half_t, cutlass::layout::RowMajor, 8, + cutlass::epilogue::TmaWarpSpecialized2Sm, + cutlass::epilogue::fusion::LinearCombination< + cutlass::half_t, + float, + void, + float + > + >::CollectiveOp; + + using CollectiveMainloop = + typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + cutlass::half_t, cutlass::layout::RowMajor, 8, + cutlass::half_t, cutlass::layout::ColumnMajor, 8, + float, + cute::Shape, + cute::Shape, + cutlass::gemm::collective::StageCountAutoCarveout(sizeof(typename CollectiveEpilogue::SharedStorage))>, + cutlass::gemm::KernelTmaWarpSpecialized2SmSm100 + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + cute::Shape, + CollectiveMainloop, + CollectiveEpilogue, + void>; + + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; +} + +namespace cutlass3x_sm100_tensorop_gemm_f16_f16_f32_void_f16_256x256x64_2x1x1_0_tnt_align8_2sm_stage6 { + using CollectiveEpilogue = + typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + cute::Shape, + cute::Shape, + cutlass::epilogue::collective::EpilogueTileAuto, + float, float, + void, cutlass::layout::RowMajor, 8, + cutlass::half_t, cutlass::layout::RowMajor, 8, + cutlass::epilogue::TmaWarpSpecialized2Sm, + cutlass::epilogue::fusion::LinearCombination< + cutlass::half_t, + float, + void, + float + > + >::CollectiveOp; + + using CollectiveMainloop = + typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + cutlass::half_t, cutlass::layout::RowMajor, 8, + cutlass::half_t, cutlass::layout::ColumnMajor, 8, + float, + cute::Shape, + cute::Shape, + cutlass::gemm::collective::StageCountAutoCarveout(sizeof(typename CollectiveEpilogue::SharedStorage))>, + cutlass::gemm::KernelTmaWarpSpecialized2SmSm100 + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + cute::Shape, + CollectiveMainloop, + CollectiveEpilogue, + void>; + + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; +} + +TEST(cutlass3x_sm100_tensorop_gemm_f16_f16_f32_void_f16_128x128x64_1x1x1_0_tnt_align8_1sm, stage_check) +{ + static constexpr int Stages = cutlass3x_sm100_tensorop_gemm_f16_f16_f32_void_f16_128x128x64_1x1x1_0_tnt_align8_1sm_stage6::CollectiveMainloop::DispatchPolicy::Stages; + EXPECT_TRUE(Stages == 6); +} + +TEST(cutlass3x_sm100_tensorop_gemm_f16_f16_f32_void_f16_128x256x64_1x1x1_0_tnt_align8_1sm, stage_check) +{ + static constexpr int Stages = cutlass3x_sm100_tensorop_gemm_f16_f16_f32_void_f16_128x256x64_1x1x1_0_tnt_align8_1sm_stage4::CollectiveMainloop::DispatchPolicy::Stages; + EXPECT_TRUE(Stages == 4); +} + +TEST(cutlass3x_sm100_tensorop_gemm_f16_f16_f32_void_f16_256x128x64_2x1x1_0_tnt_align8_2sm, stage_check) +{ + static constexpr int Stages = cutlass3x_sm100_tensorop_gemm_f16_f16_f32_void_f16_256x128x64_2x1x1_0_tnt_align8_2sm_stage8::CollectiveMainloop::DispatchPolicy::Stages; + EXPECT_TRUE(Stages == 8); +} + +TEST(cutlass3x_sm100_tensorop_gemm_f16_f16_f32_void_f16_256x256x64_2x1x1_0_tnt_align8_2sm, stage_check) +{ + static constexpr int Stages = cutlass3x_sm100_tensorop_gemm_f16_f16_f32_void_f16_256x256x64_2x1x1_0_tnt_align8_2sm_stage6::CollectiveMainloop::DispatchPolicy::Stages; + EXPECT_TRUE(Stages == 6); +} diff --git a/test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_f32_f32_f32_void_f32_stage.cu b/test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_f32_f32_f32_void_f32_stage.cu new file mode 100644 index 00000000..ae7681f8 --- /dev/null +++ b/test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_f32_f32_f32_void_f32_stage.cu @@ -0,0 +1,234 @@ +/*************************************************************************************************** + * 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 Test for sm100 dense gemm stage +*/ + +#include "../../../../common/cutlass_unit_test.h" +#include "cute/atom/mma_atom.hpp" +#include "cute/tensor.hpp" +#include "cutlass/arch/mma_sm100.h" +#include "cutlass/cutlass.h" +#include "cutlass/epilogue/collective/collective_builder.hpp" +#include "cutlass/epilogue/dispatch_policy.hpp" +#include "cutlass/epilogue/thread/linear_combination.h" +#include "cutlass/gemm/collective/collective_builder.hpp" +#include "cutlass/gemm/device/gemm_universal_adapter.h" +#include "cutlass/gemm/kernel/gemm_universal.hpp" +#include "cutlass/numeric_types.h" +#include "../../gemm_testbed_3x.hpp" + +using namespace cute; + +namespace cutlass3x_sm100_tensorop_gemm_f32_f32_f32_void_f32_128x128x32_1x1x1_0_tnt_align4_1sm_stage6 { + using CollectiveEpilogue = + typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + cute::Shape, + cute::Shape, + cutlass::epilogue::collective::EpilogueTileAuto, + float, float, + void, cutlass::layout::RowMajor, 4, + float, cutlass::layout::RowMajor, 4, + cutlass::epilogue::TmaWarpSpecialized1Sm, + cutlass::epilogue::fusion::LinearCombination< + float, + float, + void, + float + > + >::CollectiveOp; + + using CollectiveMainloop = + typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + float, cutlass::layout::RowMajor, 4, + float, cutlass::layout::ColumnMajor, 4, + float, + cute::Shape, + cute::Shape, + cutlass::gemm::collective::StageCountAutoCarveout(sizeof(typename CollectiveEpilogue::SharedStorage))>, + cutlass::gemm::KernelTmaWarpSpecialized1SmSm100 + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + cute::Shape, + CollectiveMainloop, + CollectiveEpilogue, + void>; + + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; +} + +namespace cutlass3x_sm100_tensorop_gemm_f32_f32_f32_void_f32_128x256x32_1x1x1_0_tnt_align4_1sm_stage4 { + using CollectiveEpilogue = + typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + cute::Shape, + cute::Shape, + cutlass::epilogue::collective::EpilogueTileAuto, + float, float, + void, cutlass::layout::RowMajor, 4, + float, cutlass::layout::RowMajor, 4, + cutlass::epilogue::TmaWarpSpecialized1Sm, + cutlass::epilogue::fusion::LinearCombination< + float, + float, + void, + float + > + >::CollectiveOp; + + using CollectiveMainloop = + typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + float, cutlass::layout::RowMajor, 4, + float, cutlass::layout::ColumnMajor, 4, + float, + cute::Shape, + cute::Shape, + cutlass::gemm::collective::StageCountAutoCarveout(sizeof(typename CollectiveEpilogue::SharedStorage))>, + cutlass::gemm::KernelTmaWarpSpecialized1SmSm100 + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + cute::Shape, + CollectiveMainloop, + CollectiveEpilogue, + void>; + + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; +} + +namespace cutlass3x_sm100_tensorop_gemm_f32_f32_f32_void_f32_256x128x32_2x1x1_0_tnt_align4_2sm_stage8 { + using CollectiveEpilogue = + typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + cute::Shape, + cute::Shape, + cutlass::epilogue::collective::EpilogueTileAuto, + float, float, + void, cutlass::layout::RowMajor, 4, + float, cutlass::layout::RowMajor, 4, + cutlass::epilogue::TmaWarpSpecialized2Sm, + cutlass::epilogue::fusion::LinearCombination< + float, + float, + void, + float + > + >::CollectiveOp; + + using CollectiveMainloop = + typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + float, cutlass::layout::RowMajor, 4, + float, cutlass::layout::ColumnMajor, 4, + float, + cute::Shape, + cute::Shape, + cutlass::gemm::collective::StageCountAutoCarveout(sizeof(typename CollectiveEpilogue::SharedStorage))>, + cutlass::gemm::KernelTmaWarpSpecialized2SmSm100 + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + cute::Shape, + CollectiveMainloop, + CollectiveEpilogue, + void>; + + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; +} + +namespace cutlass3x_sm100_tensorop_gemm_f32_f32_f32_void_f32_256x256x32_2x1x1_0_tnt_align4_2sm_stage6 { + using CollectiveEpilogue = + typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + cute::Shape, + cute::Shape, + cutlass::epilogue::collective::EpilogueTileAuto, + float, float, + void, cutlass::layout::RowMajor, 4, + float, cutlass::layout::RowMajor, 4, + cutlass::epilogue::TmaWarpSpecialized2Sm, + cutlass::epilogue::fusion::LinearCombination< + float, + float, + void, + float + > + >::CollectiveOp; + + using CollectiveMainloop = + typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + float, cutlass::layout::RowMajor, 4, + float, cutlass::layout::ColumnMajor, 4, + float, + cute::Shape, + cute::Shape, + cutlass::gemm::collective::StageCountAutoCarveout(sizeof(typename CollectiveEpilogue::SharedStorage))>, + cutlass::gemm::KernelTmaWarpSpecialized2SmSm100 + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + cute::Shape, + CollectiveMainloop, + CollectiveEpilogue, + void>; + + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; +} + +TEST(cutlass3x_sm100_tensorop_gemm_f32_f32_f32_void_f32_128x128x32_1x1x1_0_tnt_align4_1sm, stage_check) +{ + static constexpr int Stages = cutlass3x_sm100_tensorop_gemm_f32_f32_f32_void_f32_128x128x32_1x1x1_0_tnt_align4_1sm_stage6::CollectiveMainloop::DispatchPolicy::Stages; + EXPECT_TRUE(Stages == 6); +} + +TEST(cutlass3x_sm100_tensorop_gemm_f32_f32_f32_void_f32_128x256x32_1x1x1_0_tnt_align4_1sm, stage_check) +{ + static constexpr int Stages = cutlass3x_sm100_tensorop_gemm_f32_f32_f32_void_f32_128x256x32_1x1x1_0_tnt_align4_1sm_stage4::CollectiveMainloop::DispatchPolicy::Stages; + EXPECT_TRUE(Stages == 4); +} + +TEST(cutlass3x_sm100_tensorop_gemm_f32_f32_f32_void_f32_256x128x32_2x1x1_0_tnt_align4_2sm, stage_check) +{ + static constexpr int Stages = cutlass3x_sm100_tensorop_gemm_f32_f32_f32_void_f32_256x128x32_2x1x1_0_tnt_align4_2sm_stage8::CollectiveMainloop::DispatchPolicy::Stages; + EXPECT_TRUE(Stages == 8); +} + +TEST(cutlass3x_sm100_tensorop_gemm_f32_f32_f32_void_f32_256x256x32_2x1x1_0_tnt_align4_2sm, stage_check) +{ + static constexpr int Stages = cutlass3x_sm100_tensorop_gemm_f32_f32_f32_void_f32_256x256x32_2x1x1_0_tnt_align4_2sm_stage6::CollectiveMainloop::DispatchPolicy::Stages; + EXPECT_TRUE(Stages == 6); +} diff --git a/test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_f4_f4_f32_tensor_op_f32_runtime_datatype.cu b/test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_f4_f4_f32_tensor_op_f32_runtime_datatype.cu new file mode 100644 index 00000000..8397d294 --- /dev/null +++ b/test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_f4_f4_f32_tensor_op_f32_runtime_datatype.cu @@ -0,0 +1,154 @@ +/*************************************************************************************************** + * 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 Tests for device-wide GEMM interface +*/ + +#include + +#include "cutlass/cutlass.h" +#include "cute/tensor.hpp" +#include "cute/atom/mma_atom.hpp" + +#include "cutlass/numeric_types.h" + +#include "cutlass/gemm/device/gemm_universal_adapter.h" +#include "cutlass/gemm/kernel/gemm_universal.hpp" +#include "cutlass/gemm/collective/collective_builder.hpp" + +#include "cutlass/epilogue/dispatch_policy.hpp" +#include "cutlass/epilogue/collective/collective_builder.hpp" + +#include "cutlass/epilogue/thread/activation.h" +#include "../../../../common/cutlass_unit_test.h" + +#include "../../gemm_testbed_3x.hpp" + +using namespace cute; + +#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) + +TEST(SM100_Device_Gemm_e2m1t_e2m1n_f32t_tensorop_2sm_f32_runtime_datatype, 512x512x128_4x4x1) { + using CollectiveEpilogue = + typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + cute::Shape, + cute::Shape, + cutlass::epilogue::collective::EpilogueTileAuto, + float, float, + float, cutlass::layout::RowMajor, 4, + float, cutlass::layout::RowMajor, 4, + cutlass::epilogue::TmaWarpSpecialized1Sm, + + cutlass::epilogue::fusion::LinearCombination< + float, + float, + float, + float + > + + >::CollectiveOp; + + using CollectiveMainloop = + typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + cutlass::type_erased_dynamic_float4_t, cutlass::layout::RowMajor, 128, + cutlass::type_erased_dynamic_float4_t, cutlass::layout::ColumnMajor, 128, + float, + cute::Shape, + cute::Shape, + cutlass::gemm::collective::StageCountAutoCarveout, + cutlass::gemm::KernelTmaWarpSpecialized2SmSm100 + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + cute::Shape, + CollectiveMainloop, + CollectiveEpilogue, + void>; + + using namespace test::gemm::device; + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + + auto pass = TestRuntimeDataTypeSmall(cute::UMMA::MXF8F6F4Format::E2M1, cute::UMMA::MXF8F6F4Format::E2M1); + EXPECT_TRUE(pass); + +} + + +TEST(SM100_Device_Gemm_e2m1t_e2m1n_f32t_tensorop_1sm_f32_runtime_datatype, 256x256x128_2x2x1) { + using CollectiveEpilogue = + typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + cute::Shape, + cute::Shape, + cutlass::epilogue::collective::EpilogueTileAuto, + float, float, + float, cutlass::layout::RowMajor, 4, + float, cutlass::layout::RowMajor, 4, + cutlass::epilogue::TmaWarpSpecialized1Sm, + + cutlass::epilogue::fusion::LinearCombination< + float, + float, + float, + float + > + + >::CollectiveOp; + + using CollectiveMainloop = + typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + cutlass::type_erased_dynamic_float4_t, cutlass::layout::RowMajor, 128, + cutlass::type_erased_dynamic_float4_t, cutlass::layout::ColumnMajor, 128, + float, + cute::Shape, + cute::Shape, + cutlass::gemm::collective::StageCountAutoCarveout, + cutlass::gemm::KernelTmaWarpSpecialized1SmSm100 + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + cute::Shape, + CollectiveMainloop, + CollectiveEpilogue, + void>; + + using namespace test::gemm::device; + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + + auto pass = TestRuntimeDataTypeSmall(cute::UMMA::MXF8F6F4Format::E2M1, cute::UMMA::MXF8F6F4Format::E2M1); + EXPECT_TRUE(pass); +} + +#endif // defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) diff --git a/test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_f6_f6_f32_tensor_op_f32_runtime_datatype.cu b/test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_f6_f6_f32_tensor_op_f32_runtime_datatype.cu new file mode 100644 index 00000000..af8316fe --- /dev/null +++ b/test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_f6_f6_f32_tensor_op_f32_runtime_datatype.cu @@ -0,0 +1,154 @@ +/*************************************************************************************************** + * 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 Tests for device-wide GEMM interface +*/ + +#include + +#include "cutlass/cutlass.h" +#include "cute/tensor.hpp" +#include "cute/atom/mma_atom.hpp" + +#include "cutlass/numeric_types.h" + +#include "cutlass/gemm/device/gemm_universal_adapter.h" +#include "cutlass/gemm/kernel/gemm_universal.hpp" +#include "cutlass/gemm/collective/collective_builder.hpp" + +#include "cutlass/epilogue/dispatch_policy.hpp" +#include "cutlass/epilogue/collective/collective_builder.hpp" + +#include "cutlass/epilogue/thread/activation.h" +#include "../../../../common/cutlass_unit_test.h" + +#include "../../gemm_testbed_3x.hpp" + +using namespace cute; + +#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) + +TEST(SM100_Device_Gemm_e3m2t_e2m3n_f32t_tensorop_1sm_f32_runtime_datatype, 256x256x128_2x2x1) { + using CollectiveEpilogue = + typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + cute::Shape, + cute::Shape, + cutlass::epilogue::collective::EpilogueTileAuto, + float, float, + float, cutlass::layout::RowMajor, 4, + float, cutlass::layout::RowMajor, 4, + cutlass::epilogue::TmaWarpSpecialized1Sm, + + cutlass::epilogue::fusion::LinearCombination< + float, + float, + float, + float + > + + >::CollectiveOp; + + using CollectiveMainloop = + typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + cutlass::type_erased_dynamic_float6_t, cutlass::layout::RowMajor, 128, + cutlass::type_erased_dynamic_float6_t, cutlass::layout::ColumnMajor, 128, + float, + cute::Shape, + cute::Shape, + cutlass::gemm::collective::StageCountAutoCarveout, + cutlass::gemm::KernelTmaWarpSpecialized1SmSm100 + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + cute::Shape, + CollectiveMainloop, + CollectiveEpilogue, + void>; + + using namespace test::gemm::device; + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + + auto pass = TestRuntimeDataTypeSmall(cute::UMMA::MXF8F6F4Format::E3M2, cute::UMMA::MXF8F6F4Format::E2M3); + EXPECT_TRUE(pass); + +} + +TEST(SM100_Device_Gemm_e3m2t_e2m3n_f32t_tensorop_1sm_f32_runtime_datatype, 512x512x128_4x4x1) { + using CollectiveEpilogue = + typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + cute::Shape, + cute::Shape, + cutlass::epilogue::collective::EpilogueTileAuto, + float, float, + float, cutlass::layout::RowMajor, 4, + float, cutlass::layout::RowMajor, 4, + cutlass::epilogue::TmaWarpSpecialized1Sm, + + cutlass::epilogue::fusion::LinearCombination< + float, + float, + float, + float + > + + >::CollectiveOp; + + using CollectiveMainloop = + typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + cutlass::type_erased_dynamic_float6_t, cutlass::layout::RowMajor, 128, + cutlass::type_erased_dynamic_float6_t, cutlass::layout::ColumnMajor, 128, + float, + cute::Shape, + cute::Shape, + cutlass::gemm::collective::StageCountAutoCarveout, + cutlass::gemm::KernelTmaWarpSpecialized1SmSm100 + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + cute::Shape, + CollectiveMainloop, + CollectiveEpilogue, + void>; + + using namespace test::gemm::device; + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + + auto pass = TestRuntimeDataTypeSmall(cute::UMMA::MXF8F6F4Format::E3M2, cute::UMMA::MXF8F6F4Format::E2M3); + EXPECT_TRUE(pass); + +} + +#endif // #if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) diff --git a/test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_f8_f4_f32_tensor_op_f32_runtime_datatype.cu b/test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_f8_f4_f32_tensor_op_f32_runtime_datatype.cu new file mode 100644 index 00000000..97e3442c --- /dev/null +++ b/test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_f8_f4_f32_tensor_op_f32_runtime_datatype.cu @@ -0,0 +1,107 @@ +/*************************************************************************************************** + * 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 Tests for device-wide GEMM interface +*/ + +#include + +#include "cutlass/cutlass.h" +#include "cute/tensor.hpp" +#include "cute/atom/mma_atom.hpp" + +#include "cutlass/numeric_types.h" + +#include "cutlass/gemm/device/gemm_universal_adapter.h" +#include "cutlass/gemm/kernel/gemm_universal.hpp" +#include "cutlass/gemm/collective/collective_builder.hpp" + +#include "cutlass/epilogue/dispatch_policy.hpp" +#include "cutlass/epilogue/collective/collective_builder.hpp" + +#include "cutlass/epilogue/thread/activation.h" +#include "../../../../common/cutlass_unit_test.h" + +#include "../../gemm_testbed_3x.hpp" + +using namespace cute; + +#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) + +TEST(SM100_Device_Gemm_e4m3t_e2m1n_f32t_tensorop_2sm_f32_runtime_datatype, 256x128x128_2x2x1) { + using CollectiveEpilogue = + typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + cute::Shape, + cute::Shape, + cutlass::epilogue::collective::EpilogueTileAuto, + float, float, + float, cutlass::layout::RowMajor, 4, + float, cutlass::layout::RowMajor, 4, + cutlass::epilogue::TmaWarpSpecialized2Sm, + + cutlass::epilogue::fusion::LinearCombination< + float, + float, + float, + float + > + + >::CollectiveOp; + + using CollectiveMainloop = + typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + cutlass::type_erased_dynamic_float8_t, cutlass::layout::RowMajor, 16, + cutlass::type_erased_dynamic_float4_t, cutlass::layout::ColumnMajor, 128, + float, + cute::Shape, + cute::Shape, + cutlass::gemm::collective::StageCountAutoCarveout, + cutlass::gemm::KernelTmaWarpSpecialized2SmSm100 + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + cute::Shape, + CollectiveMainloop, + CollectiveEpilogue, + void>; + + using namespace test::gemm::device; + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + + auto pass = TestRuntimeDataTypeSmall(cute::UMMA::MXF8F6F4Format::E4M3, cute::UMMA::MXF8F6F4Format::E2M1); + EXPECT_TRUE(pass); + +} + +#endif // #if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) diff --git a/test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_f8_f8_f32_void_f8_stage.cu b/test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_f8_f8_f32_void_f8_stage.cu new file mode 100644 index 00000000..7f091314 --- /dev/null +++ b/test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_f8_f8_f32_void_f8_stage.cu @@ -0,0 +1,234 @@ +/*************************************************************************************************** + * 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 Test for sm100 gemm stage +*/ + +#include "../../../../common/cutlass_unit_test.h" +#include "cute/atom/mma_atom.hpp" +#include "cute/tensor.hpp" +#include "cutlass/arch/mma_sm100.h" +#include "cutlass/cutlass.h" +#include "cutlass/epilogue/collective/collective_builder.hpp" +#include "cutlass/epilogue/dispatch_policy.hpp" +#include "cutlass/epilogue/thread/linear_combination.h" +#include "cutlass/gemm/collective/collective_builder.hpp" +#include "cutlass/gemm/device/gemm_universal_adapter.h" +#include "cutlass/gemm/kernel/gemm_universal.hpp" +#include "cutlass/numeric_types.h" +#include "../../gemm_testbed_3x.hpp" + +using namespace cute; + +namespace cutlass3x_sm100_tensorop_gemm_e4m3_e4m3_f32_void_e4m3_128x128x128_1x1x1_0_tnt_align16_1sm_stage6 { + using CollectiveEpilogue = + typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + cute::Shape, + cute::Shape, + cutlass::epilogue::collective::EpilogueTileAuto, + float, float, + void, cutlass::layout::RowMajor, 16, + cutlass::float_e4m3_t, cutlass::layout::RowMajor, 16, + cutlass::epilogue::TmaWarpSpecialized1Sm, + cutlass::epilogue::fusion::LinearCombination< + cutlass::float_e4m3_t, + float, + void, + float + > + >::CollectiveOp; + + using CollectiveMainloop = + typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + cutlass::float_e4m3_t, cutlass::layout::RowMajor, 16, + cutlass::float_e4m3_t, cutlass::layout::ColumnMajor, 16, + float, + cute::Shape, + cute::Shape, + cutlass::gemm::collective::StageCountAutoCarveout(sizeof(typename CollectiveEpilogue::SharedStorage))>, + cutlass::gemm::KernelTmaWarpSpecialized1SmSm100 + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + cute::Shape, + CollectiveMainloop, + CollectiveEpilogue, + void>; + + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; +} + +namespace cutlass3x_sm100_tensorop_gemm_e4m3_e4m3_f32_void_e4m3_128x256x128_1x1x1_0_tnt_align16_1sm_stage4 { + using CollectiveEpilogue = + typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + cute::Shape, + cute::Shape, + cutlass::epilogue::collective::EpilogueTileAuto, + float, float, + void, cutlass::layout::RowMajor, 16, + cutlass::float_e4m3_t, cutlass::layout::RowMajor, 16, + cutlass::epilogue::TmaWarpSpecialized1Sm, + cutlass::epilogue::fusion::LinearCombination< + cutlass::float_e4m3_t, + float, + void, + float + > + >::CollectiveOp; + + using CollectiveMainloop = + typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + cutlass::float_e4m3_t, cutlass::layout::RowMajor, 16, + cutlass::float_e4m3_t, cutlass::layout::ColumnMajor, 16, + float, + cute::Shape, + cute::Shape, + cutlass::gemm::collective::StageCountAutoCarveoutEpi, + cutlass::gemm::KernelTmaWarpSpecialized1SmSm100 + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + cute::Shape, + CollectiveMainloop, + CollectiveEpilogue, + void>; + + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; +} + +namespace cutlass3x_sm100_tensorop_gemm_e4m3_e4m3_f32_void_e4m3_256x128x128_2x1x1_0_tnt_align16_2sm_stage9 { + using CollectiveEpilogue = + typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + cute::Shape, + cute::Shape, + cutlass::epilogue::collective::EpilogueTileAuto, + float, float, + void, cutlass::layout::RowMajor, 16, + cutlass::float_e4m3_t, cutlass::layout::RowMajor, 16, + cutlass::epilogue::TmaWarpSpecialized2Sm, + cutlass::epilogue::fusion::LinearCombination< + cutlass::float_e4m3_t, + float, + void, + float + > + >::CollectiveOp; + + using CollectiveMainloop = + typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + cutlass::float_e4m3_t, cutlass::layout::RowMajor, 16, + cutlass::float_e4m3_t, cutlass::layout::ColumnMajor, 16, + float, + cute::Shape, + cute::Shape, + cutlass::gemm::collective::StageCountAutoCarveout(sizeof(typename CollectiveEpilogue::SharedStorage))>, + cutlass::gemm::KernelTmaWarpSpecialized2SmSm100 + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + cute::Shape, + CollectiveMainloop, + CollectiveEpilogue, + void>; + + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; +} + +namespace cutlass3x_sm100_tensorop_gemm_e4m3_e4m3_f32_void_e4m3_256x256x128_2x1x1_0_tnt_align16_2sm_stage6 { + using CollectiveEpilogue = + typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + cute::Shape, + cute::Shape, + cutlass::epilogue::collective::EpilogueTileAuto, + float, float, + void, cutlass::layout::RowMajor, 16, + cutlass::float_e4m3_t, cutlass::layout::RowMajor, 16, + cutlass::epilogue::TmaWarpSpecialized2Sm, + cutlass::epilogue::fusion::LinearCombination< + cutlass::float_e4m3_t, + float, + void, + float + > + >::CollectiveOp; + + using CollectiveMainloop = + typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + cutlass::float_e4m3_t, cutlass::layout::RowMajor, 16, + cutlass::float_e4m3_t, cutlass::layout::ColumnMajor, 16, + float, + cute::Shape, + cute::Shape, + cutlass::gemm::collective::StageCountAutoCarveout(sizeof(typename CollectiveEpilogue::SharedStorage))>, + cutlass::gemm::KernelTmaWarpSpecialized2SmSm100 + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + cute::Shape, + CollectiveMainloop, + CollectiveEpilogue, + void>; + + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; +} + +TEST(cutlass3x_sm100_tensorop_gemm_e4m3_e4m3_f32_void_e4m3_128x128x128_1x1x1_0_tnt_align16_1sm, stage_check) +{ + static constexpr int Stages = cutlass3x_sm100_tensorop_gemm_e4m3_e4m3_f32_void_e4m3_128x128x128_1x1x1_0_tnt_align16_1sm_stage6::CollectiveMainloop::DispatchPolicy::Stages; + EXPECT_TRUE(Stages == 6); +} + +TEST(cutlass3x_sm100_tensorop_gemm_e4m3_e4m3_f32_void_e4m3_128x256x128_1x1x1_0_tnt_align16_1sm, stage_check) +{ + static constexpr int Stages = cutlass3x_sm100_tensorop_gemm_e4m3_e4m3_f32_void_e4m3_128x256x128_1x1x1_0_tnt_align16_1sm_stage4::CollectiveMainloop::DispatchPolicy::Stages; + EXPECT_TRUE(Stages == 4); +} + +TEST(cutlass3x_sm100_tensorop_gemm_e4m3_e4m3_f32_void_e4m3_256x128x128_2x1x1_0_tnt_align16_2sm, stage_check) +{ + static constexpr int Stages = cutlass3x_sm100_tensorop_gemm_e4m3_e4m3_f32_void_e4m3_256x128x128_2x1x1_0_tnt_align16_2sm_stage9::CollectiveMainloop::DispatchPolicy::Stages; + EXPECT_TRUE(Stages == 9); +} + +TEST(cutlass3x_sm100_tensorop_gemm_e4m3_e4m3_f32_void_e4m3_256x256x128_2x1x1_0_tnt_align16_2sm, stage_check) +{ + static constexpr int Stages = cutlass3x_sm100_tensorop_gemm_e4m3_e4m3_f32_void_e4m3_256x256x128_2x1x1_0_tnt_align16_2sm_stage6::CollectiveMainloop::DispatchPolicy::Stages; + EXPECT_TRUE(Stages == 6); +} diff --git a/test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_f8_f8_f8_tensor_op_f32_bias_gelu.cu b/test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_f8_f8_f8_tensor_op_f32_bias_gelu.cu new file mode 100644 index 00000000..17428cb2 --- /dev/null +++ b/test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_f8_f8_f8_tensor_op_f32_bias_gelu.cu @@ -0,0 +1,319 @@ +/*************************************************************************************************** + * 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 Tests for device-wide GEMM interface +*/ + +#include + +#include "cutlass/cutlass.h" +#include "cute/tensor.hpp" +#include "cute/atom/mma_atom.hpp" + +#include "cutlass/numeric_types.h" +#include "cutlass/arch/mma_sm100.h" + +#include "cutlass/gemm/device/gemm_universal_adapter.h" +#include "cutlass/gemm/kernel/gemm_universal.hpp" +#include "cutlass/gemm/collective/collective_builder.hpp" + +#include "cutlass/epilogue/dispatch_policy.hpp" +#include "cutlass/epilogue/collective/collective_builder.hpp" + +#include "cutlass/epilogue/thread/activation.h" +#include "../../../../common/cutlass_unit_test.h" + +#include "../../gemm_testbed_3x.hpp" + +using namespace cute; + +#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) + +/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// +///////////////////////////////////////////////////// 128x128x128 ////////////////////////////////////////////////////// +/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// + +TEST(SM100_Device_Gemm_e4m3t_e4m3n_e4m3n_tensorop_1sm_f32_bias_gelu, 128x128x128_1x1x1) { + using LayoutA = cutlass::layout::RowMajor; + using LayoutB = cutlass::layout::ColumnMajor; + using LayoutC = cutlass::layout::ColumnMajor; + using ElementA = cutlass::float_e4m3_t; + using ElementB = cutlass::float_e4m3_t; + using ElementC = cutlass::float_e4m3_t; + using ElementD = cutlass::float_e4m3_t; + using ElementAccumulator = float; + using ElementCompute = float; + using ElementBias = cutlass::half_t; + using MmaTileShape = cute::Shape<_128,_128,Int<128 / sizeof(ElementA)>>; + using ClusterShape = Shape<_1,_1,_1>; + + using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecialized1Sm; + using FusionOperation = cutlass::epilogue::fusion::ScaledLinCombPerRowBiasEltAct< + cutlass::epilogue::thread::ScaledGELU_taylor, ElementD, ElementCompute, ElementBias>; + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + MmaTileShape, ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + ElementAccumulator, ElementCompute, + ElementC, LayoutC, 16 / sizeof(ElementC), + ElementD, LayoutC, 16 / sizeof(ElementD), + EpilogueSchedule, + FusionOperation + >::CollectiveOp; + + using MainloopSchedule = cutlass::gemm::KernelTmaWarpSpecialized1SmSm100; + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + ElementA, LayoutA, 16 / sizeof(ElementA), + ElementB, LayoutB, 16 / sizeof(ElementB), + ElementAccumulator, + MmaTileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout(sizeof(typename CollectiveEpilogue::SharedStorage))>, + MainloopSchedule + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, + CollectiveMainloop, + CollectiveEpilogue + >; + + using namespace test::gemm::device; + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + auto pass = TestSmallFusion(1.0, 0.5, CheckEquality::RELATIVE); + EXPECT_TRUE(pass); +} + +TEST(SM100_Device_Gemm_e4m3t_e4m3n_e4m3n_tensorop_1sm_f32_bias_gelu, 64x128x128_1x2x1) { + using LayoutA = cutlass::layout::RowMajor; + using LayoutB = cutlass::layout::ColumnMajor; + using LayoutC = cutlass::layout::ColumnMajor; + using ElementA = cutlass::float_e4m3_t; + using ElementB = cutlass::float_e4m3_t; + using ElementC = cutlass::float_e4m3_t; + using ElementD = cutlass::float_e4m3_t; + using ElementAccumulator = float; + using ElementCompute = float; + using ElementBias = cutlass::half_t; + using MmaTileShape = cute::Shape<_64,_64,Int<128 / sizeof(ElementA)>>; + using ClusterShape = Shape<_1,_2,_1>; + + using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecialized1Sm; + using FusionOperation = cutlass::epilogue::fusion::ScaledLinCombPerRowBiasEltAct< + cutlass::epilogue::thread::ScaledGELU_taylor, ElementD, ElementCompute, ElementBias>; + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + MmaTileShape, ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + ElementAccumulator, ElementCompute, + ElementC, LayoutC, 16 / sizeof(ElementC), + ElementD, LayoutC, 16 / sizeof(ElementD), + EpilogueSchedule, + FusionOperation + >::CollectiveOp; + + using MainloopSchedule = cutlass::gemm::KernelTmaWarpSpecialized1SmSm100; + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + ElementA, LayoutA, 16 / sizeof(ElementA), + ElementB, LayoutB, 16 / sizeof(ElementB), + ElementAccumulator, + MmaTileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout(sizeof(typename CollectiveEpilogue::SharedStorage))>, + MainloopSchedule + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, + CollectiveMainloop, + CollectiveEpilogue + >; + + using namespace test::gemm::device; + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + auto pass = TestSmallFusion(1.0, 0.5, CheckEquality::RELATIVE); + EXPECT_TRUE(pass); +} + +TEST(SM100_Device_Gemm_e4m3t_e4m3n_e4m3n_tensorop_2sm_f32_bias_gelu, 256x128x128_2x1x1) { + using LayoutA = cutlass::layout::RowMajor; + using LayoutB = cutlass::layout::ColumnMajor; + using LayoutC = cutlass::layout::ColumnMajor; + using ElementA = cutlass::float_e4m3_t; + using ElementB = cutlass::float_e4m3_t; + using ElementC = cutlass::float_e4m3_t; + using ElementD = cutlass::float_e4m3_t; + using ElementAccumulator = float; + using ElementCompute = float; + using ElementBias = cutlass::half_t; + using MmaTileShape = cute::Shape<_256,_128,Int<128 / sizeof(ElementA)>>; + using ClusterShape = Shape<_2,_1,_1>; + + using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecialized2Sm; + using FusionOperation = cutlass::epilogue::fusion::ScaledLinCombPerRowBiasEltAct< + cutlass::epilogue::thread::ScaledGELU_taylor, ElementD, ElementCompute, ElementBias>; + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + MmaTileShape, ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + ElementAccumulator, ElementCompute, + ElementC, LayoutC, 16 / sizeof(ElementC), + ElementD, LayoutC, 16 / sizeof(ElementD), + EpilogueSchedule, + FusionOperation + >::CollectiveOp; + + using MainloopSchedule = cutlass::gemm::KernelTmaWarpSpecialized2SmSm100; + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + ElementA, LayoutA, 16 / sizeof(ElementA), + ElementB, LayoutB, 16 / sizeof(ElementB), + ElementAccumulator, + MmaTileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout(sizeof(typename CollectiveEpilogue::SharedStorage))>, + MainloopSchedule + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, + CollectiveMainloop, + CollectiveEpilogue + >; + + using namespace test::gemm::device; + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + auto pass = TestSmallFusion(1.0, 0.5, CheckEquality::RELATIVE); + EXPECT_TRUE(pass); +} + +TEST(SM100_Device_Gemm_e4m3t_e4m3n_e4m3n_tensorop_2sm_f32_bias_gelu, 512x512x128_4x4x1) { + using LayoutA = cutlass::layout::RowMajor; + using LayoutB = cutlass::layout::ColumnMajor; + using LayoutC = cutlass::layout::ColumnMajor; + using ElementA = cutlass::float_e4m3_t; + using ElementB = cutlass::float_e4m3_t; + using ElementC = cutlass::float_e4m3_t; + using ElementD = cutlass::float_e4m3_t; + using ElementAccumulator = float; + using ElementCompute = float; + using ElementBias = cutlass::half_t; + using MmaTileShape = cute::Shape<_256,_128,Int<128 / sizeof(ElementA)>>; + using ClusterShape = Shape<_4,_4,_1>; + + using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecialized2Sm; + using FusionOperation = cutlass::epilogue::fusion::ScaledLinCombPerRowBiasEltAct< + cutlass::epilogue::thread::ScaledGELU_taylor, ElementD, ElementCompute, ElementBias>; + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + MmaTileShape, ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + ElementAccumulator, ElementCompute, + ElementC, LayoutC, 16 / sizeof(ElementC), + ElementD, LayoutC, 16 / sizeof(ElementD), + EpilogueSchedule, + FusionOperation + >::CollectiveOp; + + using MainloopSchedule = cutlass::gemm::KernelTmaWarpSpecialized2SmSm100; + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + ElementA, LayoutA, 16 / sizeof(ElementA), + ElementB, LayoutB, 16 / sizeof(ElementB), + ElementAccumulator, + MmaTileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout(sizeof(typename CollectiveEpilogue::SharedStorage))>, + MainloopSchedule + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, + CollectiveMainloop, + CollectiveEpilogue + >; + + using namespace test::gemm::device; + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + auto pass = TestSmallFusion(1.0, 0.5, CheckEquality::RELATIVE); + EXPECT_TRUE(pass); +} + +TEST(SM100_Device_Gemm_e4m3t_e4m3n_e4m3n_tensorop_2sm_f32_bias_dgelu, 512x512x128_4x4x1) { + using LayoutA = cutlass::layout::RowMajor; + using LayoutB = cutlass::layout::ColumnMajor; + using LayoutC = cutlass::layout::ColumnMajor; + using ElementA = cutlass::float_e4m3_t; + using ElementB = cutlass::float_e4m3_t; + using ElementC = cutlass::float_e4m3_t; + using ElementD = cutlass::float_e4m3_t; + using ElementAccumulator = float; + using ElementCompute = float; + using ElementBias = cutlass::half_t; + using MmaTileShape = cute::Shape<_256,_128,Int<128 / sizeof(ElementA)>>; + using ClusterShape = Shape<_4,_4,_1>; + + using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecialized2Sm; + using FusionOperation = cutlass::epilogue::fusion::LinCombDeEltAct< + LayoutC, cutlass::epilogue::thread::dGELU, ElementD, ElementCompute, ElementD>; + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + MmaTileShape, ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + ElementAccumulator, ElementCompute, + ElementC, LayoutC, 16 / sizeof(ElementC), + ElementD, LayoutC, 16 / sizeof(ElementD), + EpilogueSchedule, + FusionOperation + >::CollectiveOp; + + using MainloopSchedule = cutlass::gemm::KernelTmaWarpSpecialized2SmSm100; + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + ElementA, LayoutA, 16 / sizeof(ElementA), + ElementB, LayoutB, 16 / sizeof(ElementB), + ElementAccumulator, + MmaTileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout(sizeof(typename CollectiveEpilogue::SharedStorage))>, + MainloopSchedule + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, + CollectiveMainloop, + CollectiveEpilogue + >; + + using namespace test::gemm::device; + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + auto pass = TestSmallFusion(1.0, 0.5, CheckEquality::RELATIVE); + EXPECT_TRUE(pass); +} +#endif // #if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) diff --git a/test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_f8_f8_f8_tensor_op_f32_bias_gelu_amax_aux.cu b/test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_f8_f8_f8_tensor_op_f32_bias_gelu_amax_aux.cu new file mode 100644 index 00000000..c801f6d9 --- /dev/null +++ b/test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_f8_f8_f8_tensor_op_f32_bias_gelu_amax_aux.cu @@ -0,0 +1,328 @@ +/*************************************************************************************************** + * 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 Tests for device-wide GEMM interface +*/ + +#include + +#include "cutlass/cutlass.h" +#include "cute/tensor.hpp" +#include "cute/atom/mma_atom.hpp" + +#include "cutlass/numeric_types.h" +#include "cutlass/arch/mma_sm100.h" + +#include "cutlass/gemm/device/gemm_universal_adapter.h" +#include "cutlass/gemm/kernel/gemm_universal.hpp" +#include "cutlass/gemm/collective/collective_builder.hpp" + +#include "cutlass/epilogue/dispatch_policy.hpp" +#include "cutlass/epilogue/collective/collective_builder.hpp" + +#include "cutlass/epilogue/thread/activation.h" +#include "../../../../common/cutlass_unit_test.h" + +#include "../../gemm_testbed_3x.hpp" + +using namespace cute; + +#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) + +/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// +///////////////////////////////////////////////////// 128x128x128 ////////////////////////////////////////////////////// +/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// + +TEST(SM100_Device_Gemm_e4m3t_e4m3n_e4m3n_tensorop_1sm_f32_bias_relu_amax_aux, 128x128x128_1x1x1) { + using LayoutA = cutlass::layout::RowMajor; + using LayoutB = cutlass::layout::ColumnMajor; + using LayoutC = cutlass::layout::ColumnMajor; + using ElementA = cutlass::float_e4m3_t; + using ElementB = cutlass::float_e4m3_t; + using ElementC = cutlass::float_e4m3_t; + using ElementD = cutlass::float_e4m3_t; + using ElementAccumulator = float; + using ElementCompute = float; + using ElementBias = cutlass::half_t; + using ElementAmax = float; + using ElementAux = cutlass::uint1b_t; + using MmaTileShape = cute::Shape<_128,_128,Int<128 / sizeof(ElementA)>>; + using ClusterShape = Shape<_1,_1,_1>; + + using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecialized1Sm; + using FusionOperation = cutlass::epilogue::fusion::ScaledLinCombPerRowBiasEltActAmaxAux< + LayoutC, cutlass::epilogue::thread::ReLU, ElementD, ElementCompute, ElementAux, ElementAmax, ElementBias>; + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + MmaTileShape, ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + ElementAccumulator, ElementCompute, + ElementC, LayoutC, 16 / sizeof(ElementC), + ElementD, LayoutC, 16 / sizeof(ElementD), + EpilogueSchedule, + FusionOperation + >::CollectiveOp; + + using MainloopSchedule = cutlass::gemm::KernelTmaWarpSpecialized1SmSm100; + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + ElementA, LayoutA, 16 / sizeof(ElementA), + ElementB, LayoutB, 16 / sizeof(ElementB), + ElementAccumulator, + MmaTileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout(sizeof(typename CollectiveEpilogue::SharedStorage))>, + MainloopSchedule + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, + CollectiveMainloop, + CollectiveEpilogue + >; + using namespace test::gemm::device; + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + auto pass = TestSmallFusion(1.0, 0.5, CheckEquality::RELATIVE); + EXPECT_TRUE(pass); +} + + +TEST(SM100_Device_Gemm_e4m3t_e4m3n_e4m3t_tensorop_1sm_f32_colbias_relu_amax_aux, 128x128x128_1x1x1) { + using LayoutA = cutlass::layout::RowMajor; + using LayoutB = cutlass::layout::ColumnMajor; + using LayoutC = cutlass::layout::RowMajor; + using ElementA = cutlass::float_e4m3_t; + using ElementB = cutlass::float_e4m3_t; + using ElementC = cutlass::float_e4m3_t; + using ElementD = cutlass::float_e4m3_t; + using ElementAccumulator = float; + using ElementCompute = float; + using ElementBias = cutlass::half_t; + using ElementAmax = float; + using ElementAux = cutlass::uint1b_t; + using MmaTileShape = cute::Shape<_128,_128,Int<128 / sizeof(ElementA)>>; + using ClusterShape = Shape<_1,_1,_1>; + + using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecialized1Sm; + using FusionOperation = cutlass::epilogue::fusion::ScaledLinCombPerColBiasEltActAmaxAux< + LayoutC, cutlass::epilogue::thread::ReLU, ElementD, ElementCompute, ElementAux, ElementAmax, ElementBias>; + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + MmaTileShape, ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + ElementAccumulator, ElementCompute, + ElementC, LayoutC, 16 / sizeof(ElementC), + ElementD, LayoutC, 16 / sizeof(ElementD), + EpilogueSchedule, + FusionOperation + >::CollectiveOp; + + using MainloopSchedule = cutlass::gemm::KernelTmaWarpSpecialized1SmSm100; + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + ElementA, LayoutA, 16 / sizeof(ElementA), + ElementB, LayoutB, 16 / sizeof(ElementB), + ElementAccumulator, + MmaTileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout(sizeof(typename CollectiveEpilogue::SharedStorage))>, + MainloopSchedule + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, + CollectiveMainloop, + CollectiveEpilogue + >; + using namespace test::gemm::device; + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + auto pass = TestSmallFusion(1.0, 0.5, CheckEquality::RELATIVE); + EXPECT_TRUE(pass); +} + +TEST(SM100_Device_Gemm_e4m3t_e4m3n_e4m3n_tensorop_1sm_f32_bias_gelu_amax_aux, 64x128x128_1x2x1) { + using LayoutA = cutlass::layout::RowMajor; + using LayoutB = cutlass::layout::ColumnMajor; + using LayoutC = cutlass::layout::ColumnMajor; + using ElementA = cutlass::float_e4m3_t; + using ElementB = cutlass::float_e4m3_t; + using ElementC = cutlass::float_e4m3_t; + using ElementD = cutlass::float_e4m3_t; + using ElementAccumulator = float; + using ElementCompute = float; + using ElementBias = cutlass::half_t; + using ElementAmax = float; + using ElementAux = cutlass::float_e4m3_t; + using MmaTileShape = cute::Shape<_64,_64,Int<128 / sizeof(ElementA)>>; + using ClusterShape = Shape<_1,_2,_1>; + + using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecialized1Sm; + using FusionOperation = cutlass::epilogue::fusion::ScaledLinCombPerRowBiasEltActAmaxAux< + LayoutC, cutlass::epilogue::thread::GELU, ElementD, ElementCompute, ElementAux, ElementAmax, ElementBias>; + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + MmaTileShape, ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + ElementAccumulator, ElementCompute, + ElementC, LayoutC, 16 / sizeof(ElementC), + ElementD, LayoutC, 16 / sizeof(ElementD), + EpilogueSchedule, + FusionOperation + >::CollectiveOp; + + using MainloopSchedule = cutlass::gemm::KernelTmaWarpSpecialized1SmSm100; + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + ElementA, LayoutA, 16 / sizeof(ElementA), + ElementB, LayoutB, 16 / sizeof(ElementB), + ElementAccumulator, + MmaTileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout(sizeof(typename CollectiveEpilogue::SharedStorage))>, + MainloopSchedule + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, + CollectiveMainloop, + CollectiveEpilogue + >; + + using namespace test::gemm::device; + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + auto pass = TestSmallFusion(1.0, 0.5, CheckEquality::RELATIVE); + EXPECT_TRUE(pass); +} + +TEST(SM100_Device_Gemm_e4m3t_e4m3n_e4m3n_tensorop_2sm_f32_bias_gelu_amax_aux, 256x128x128_2x1x1) { + using LayoutA = cutlass::layout::RowMajor; + using LayoutB = cutlass::layout::ColumnMajor; + using LayoutC = cutlass::layout::ColumnMajor; + using ElementA = cutlass::float_e4m3_t; + using ElementB = cutlass::float_e4m3_t; + using ElementC = cutlass::float_e4m3_t; + using ElementD = cutlass::float_e4m3_t; + using ElementAccumulator = float; + using ElementCompute = float; + using ElementBias = cutlass::half_t; + using ElementAmax = float; + using ElementAux = cutlass::float_e4m3_t; + using MmaTileShape = cute::Shape<_256,_128,Int<128 / sizeof(ElementA)>>; + using ClusterShape = Shape<_2,_1,_1>; + + using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecialized2Sm; + using FusionOperation = cutlass::epilogue::fusion::ScaledLinCombPerRowBiasEltActAmaxAux< + LayoutC, cutlass::epilogue::thread::GELU, ElementD, ElementCompute, ElementAux, ElementAmax, ElementBias>; + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + MmaTileShape, ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + ElementAccumulator, ElementCompute, + ElementC, LayoutC, 16 / sizeof(ElementC), + ElementD, LayoutC, 16 / sizeof(ElementD), + EpilogueSchedule, + FusionOperation + >::CollectiveOp; + + using MainloopSchedule = cutlass::gemm::KernelTmaWarpSpecialized2SmSm100; + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + ElementA, LayoutA, 16 / sizeof(ElementA), + ElementB, LayoutB, 16 / sizeof(ElementB), + ElementAccumulator, + MmaTileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout(sizeof(typename CollectiveEpilogue::SharedStorage))>, + MainloopSchedule + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, + CollectiveMainloop, + CollectiveEpilogue + >; + + using namespace test::gemm::device; + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + auto pass = TestSmallFusion(1.0, 0.5, CheckEquality::RELATIVE); + EXPECT_TRUE(pass); +} + +TEST(SM100_Device_Gemm_e4m3t_e4m3n_e4m3n_tensorop_2sm_f32_bias_gelu_amax_aux, 512x512x128_4x4x1) { + using LayoutA = cutlass::layout::RowMajor; + using LayoutB = cutlass::layout::ColumnMajor; + using LayoutC = cutlass::layout::ColumnMajor; + using ElementA = cutlass::float_e4m3_t; + using ElementB = cutlass::float_e4m3_t; + using ElementC = cutlass::float_e4m3_t; + using ElementD = cutlass::float_e4m3_t; + using ElementAccumulator = float; + using ElementCompute = float; + using ElementBias = cutlass::half_t; + using ElementAmax = float; + using ElementAux = cutlass::float_e4m3_t; + using MmaTileShape = cute::Shape<_256,_128,Int<128 / sizeof(ElementA)>>; + using ClusterShape = Shape<_4,_4,_1>; + + using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecialized2Sm; + using FusionOperation = cutlass::epilogue::fusion::ScaledLinCombPerRowBiasEltActAmaxAux< + LayoutC, cutlass::epilogue::thread::GELU, ElementD, ElementCompute, ElementAux, ElementAmax, ElementBias>; + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + MmaTileShape, ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + ElementAccumulator, ElementCompute, + ElementC, LayoutC, 16 / sizeof(ElementC), + ElementD, LayoutC, 16 / sizeof(ElementD), + EpilogueSchedule, + FusionOperation + >::CollectiveOp; + + using MainloopSchedule = cutlass::gemm::KernelTmaWarpSpecialized2SmSm100; + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + ElementA, LayoutA, 16 / sizeof(ElementA), + ElementB, LayoutB, 16 / sizeof(ElementB), + ElementAccumulator, + MmaTileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout(sizeof(typename CollectiveEpilogue::SharedStorage))>, + MainloopSchedule + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, + CollectiveMainloop, + CollectiveEpilogue + >; + + using namespace test::gemm::device; + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + auto pass = TestSmallFusion(1.0, 0.5, CheckEquality::RELATIVE); + EXPECT_TRUE(pass); +} +#endif // #if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) diff --git a/test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_f8_f8_f8_tensor_op_f32_bias_relu.cu b/test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_f8_f8_f8_tensor_op_f32_bias_relu.cu new file mode 100644 index 00000000..3bdb269d --- /dev/null +++ b/test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_f8_f8_f8_tensor_op_f32_bias_relu.cu @@ -0,0 +1,365 @@ +/*************************************************************************************************** + * 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 Tests for device-wide GEMM interface +*/ + +#include + +#include "cutlass/cutlass.h" +#include "cute/tensor.hpp" +#include "cute/atom/mma_atom.hpp" + +#include "cutlass/numeric_types.h" +#include "cutlass/arch/mma_sm100.h" + +#include "cutlass/gemm/device/gemm_universal_adapter.h" +#include "cutlass/gemm/kernel/gemm_universal.hpp" +#include "cutlass/gemm/collective/collective_builder.hpp" + +#include "cutlass/epilogue/dispatch_policy.hpp" +#include "cutlass/epilogue/collective/collective_builder.hpp" + +#include "cutlass/epilogue/thread/activation.h" +#include "../../../../common/cutlass_unit_test.h" + +#include "../../gemm_testbed_3x.hpp" + +using namespace cute; + +#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) + +/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// +///////////////////////////////////////////////////// 128x128x128 ////////////////////////////////////////////////////// +/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// + +TEST(SM100_Device_Gemm_e4m3t_e4m3n_e4m3n_tensorop_1sm_f32_bias_relu, 128x128x128_1x1x1) { + using LayoutA = cutlass::layout::RowMajor; + using LayoutB = cutlass::layout::ColumnMajor; + using LayoutC = cutlass::layout::ColumnMajor; + using ElementA = cutlass::float_e4m3_t; + using ElementB = cutlass::float_e4m3_t; + using ElementC = cutlass::float_e4m3_t; + using ElementD = cutlass::float_e4m3_t; + using ElementAccumulator = float; + using ElementCompute = float; + using ElementBias = cutlass::half_t; + using MmaTileShape = cute::Shape<_128,_128,Int<128 / sizeof(ElementA)>>; + using ClusterShape = Shape<_1,_1,_1>; + + using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecialized1Sm; + using FusionOperation = cutlass::epilogue::fusion::ScaledLinCombPerRowBiasEltAct< + cutlass::epilogue::thread::ReLU, ElementD, ElementCompute, ElementBias>; + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + MmaTileShape, ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + ElementAccumulator, ElementCompute, + ElementC, LayoutC, 16 / sizeof(ElementC), + ElementD, LayoutC, 16 / sizeof(ElementD), + EpilogueSchedule, + FusionOperation + >::CollectiveOp; + + using MainloopSchedule = cutlass::gemm::KernelTmaWarpSpecialized1SmSm100; + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + ElementA, LayoutA, 16 / sizeof(ElementA), + ElementB, LayoutB, 16 / sizeof(ElementB), + ElementAccumulator, + MmaTileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout(sizeof(typename CollectiveEpilogue::SharedStorage))>, + MainloopSchedule + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, + CollectiveMainloop, + CollectiveEpilogue + >; + + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + auto pass = test::gemm::device::TestSmallFusion(1.0, 0.5); + EXPECT_TRUE(pass); +} + +TEST(SM100_Device_Gemm_e4m3t_e4m3n_e4m3n_tensorop_1sm_f32_bias_relu_beta0, 128x128x128_1x1x1) { + using LayoutA = cutlass::layout::RowMajor; + using LayoutB = cutlass::layout::ColumnMajor; + using LayoutC = cutlass::layout::ColumnMajor; + using ElementA = cutlass::float_e4m3_t; + using ElementB = cutlass::float_e4m3_t; + using ElementC = cutlass::float_e4m3_t; + using ElementD = cutlass::float_e4m3_t; + using ElementAccumulator = float; + using ElementCompute = float; + using ElementBias = cutlass::half_t; + using MmaTileShape = cute::Shape<_128,_128,Int<128 / sizeof(ElementA)>>; + using ClusterShape = Shape<_1,_1,_1>; + + using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecialized1Sm; + using FusionOperation = cutlass::epilogue::fusion::ScaledLinCombPerRowBiasEltAct< + cutlass::epilogue::thread::ReLU, ElementD, ElementCompute, ElementBias>; + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + MmaTileShape, ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + ElementAccumulator, ElementCompute, + ElementC, LayoutC, 16 / sizeof(ElementC), + ElementD, LayoutC, 16 / sizeof(ElementD), + EpilogueSchedule, + FusionOperation + >::CollectiveOp; + + using MainloopSchedule = cutlass::gemm::KernelTmaWarpSpecialized1SmSm100; + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + ElementA, LayoutA, 16 / sizeof(ElementA), + ElementB, LayoutB, 16 / sizeof(ElementB), + ElementAccumulator, + MmaTileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout(sizeof(typename CollectiveEpilogue::SharedStorage))>, + MainloopSchedule + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, + CollectiveMainloop, + CollectiveEpilogue + >; + + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + auto pass = test::gemm::device::TestSmallFusion(1.0, 0); + EXPECT_TRUE(pass); +} + + +TEST(SM100_Device_Gemm_e4m3t_e4m3n_e4m3t_tensorop_1sm_f32_colbias_relu, 128x128x128_1x1x1) { + using LayoutA = cutlass::layout::RowMajor; + using LayoutB = cutlass::layout::ColumnMajor; + using LayoutC = cutlass::layout::RowMajor; + using ElementA = cutlass::float_e4m3_t; + using ElementB = cutlass::float_e4m3_t; + using ElementC = cutlass::float_e4m3_t; + using ElementD = cutlass::float_e4m3_t; + using ElementAccumulator = float; + using ElementCompute = float; + using ElementBias = cutlass::half_t; + using MmaTileShape = cute::Shape<_128,_128,Int<128 / sizeof(ElementA)>>; + using ClusterShape = Shape<_1,_1,_1>; + + using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecialized1Sm; + using FusionOperation = cutlass::epilogue::fusion::ScaledLinCombPerColBiasEltAct< + cutlass::epilogue::thread::ReLU, ElementD, ElementCompute, ElementBias>; + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + MmaTileShape, ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + ElementAccumulator, ElementCompute, + ElementC, LayoutC, 16 / sizeof(ElementC), + ElementD, LayoutC, 16 / sizeof(ElementD), + EpilogueSchedule, + FusionOperation + >::CollectiveOp; + + using MainloopSchedule = cutlass::gemm::KernelTmaWarpSpecialized1SmSm100; + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + ElementA, LayoutA, 16 / sizeof(ElementA), + ElementB, LayoutB, 16 / sizeof(ElementB), + ElementAccumulator, + MmaTileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout(sizeof(typename CollectiveEpilogue::SharedStorage))>, + MainloopSchedule + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, + CollectiveMainloop, + CollectiveEpilogue + >; + + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + auto pass = test::gemm::device::TestSmallFusion(1.0, 0.5); + EXPECT_TRUE(pass); +} + +TEST(SM100_Device_Gemm_e4m3t_e4m3n_e4m3n_tensorop_1sm_f32_bias_relu, 64x128x128_1x2x1) { + using LayoutA = cutlass::layout::RowMajor; + using LayoutB = cutlass::layout::ColumnMajor; + using LayoutC = cutlass::layout::ColumnMajor; + using ElementA = cutlass::float_e4m3_t; + using ElementB = cutlass::float_e4m3_t; + using ElementC = cutlass::float_e4m3_t; + using ElementD = cutlass::float_e4m3_t; + using ElementAccumulator = float; + using ElementCompute = float; + using ElementBias = cutlass::half_t; + using MmaTileShape = cute::Shape<_64,_64,Int<128 / sizeof(ElementA)>>; + using ClusterShape = Shape<_1,_2,_1>; + + using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecialized1Sm; + using FusionOperation = cutlass::epilogue::fusion::ScaledLinCombPerRowBiasEltAct< + cutlass::epilogue::thread::ReLU, ElementD, ElementCompute, ElementBias>; + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + MmaTileShape, ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + ElementAccumulator, ElementCompute, + ElementC, LayoutC, 16 / sizeof(ElementC), + ElementD, LayoutC, 16 / sizeof(ElementD), + EpilogueSchedule, + FusionOperation + >::CollectiveOp; + + using MainloopSchedule = cutlass::gemm::KernelTmaWarpSpecialized1SmSm100; + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + ElementA, LayoutA, 16 / sizeof(ElementA), + ElementB, LayoutB, 16 / sizeof(ElementB), + ElementAccumulator, + MmaTileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout(sizeof(typename CollectiveEpilogue::SharedStorage))>, + MainloopSchedule + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, + CollectiveMainloop, + CollectiveEpilogue + >; + + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + auto pass = test::gemm::device::TestSmallFusion(1.0, 0.5); + EXPECT_TRUE(pass); +} + +TEST(SM100_Device_Gemm_e4m3t_e4m3n_e4m3n_tensorop_2sm_f32_bias_relu, 256x128x128_2x1x1) { + using LayoutA = cutlass::layout::RowMajor; + using LayoutB = cutlass::layout::ColumnMajor; + using LayoutC = cutlass::layout::ColumnMajor; + using ElementA = cutlass::float_e4m3_t; + using ElementB = cutlass::float_e4m3_t; + using ElementC = cutlass::float_e4m3_t; + using ElementD = cutlass::float_e4m3_t; + using ElementAccumulator = float; + using ElementCompute = float; + using ElementBias = cutlass::half_t; + using MmaTileShape = cute::Shape<_256,_128,Int<128 / sizeof(ElementA)>>; + using ClusterShape = Shape<_2,_1,_1>; + + using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecialized2Sm; + using FusionOperation = cutlass::epilogue::fusion::ScaledLinCombPerRowBiasEltAct< + cutlass::epilogue::thread::ReLU, ElementD, ElementCompute, ElementBias>; + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + MmaTileShape, ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + ElementAccumulator, ElementCompute, + ElementC, LayoutC, 16 / sizeof(ElementC), + ElementD, LayoutC, 16 / sizeof(ElementD), + EpilogueSchedule, + FusionOperation + >::CollectiveOp; + + using MainloopSchedule = cutlass::gemm::KernelTmaWarpSpecialized2SmSm100; + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + ElementA, LayoutA, 16 / sizeof(ElementA), + ElementB, LayoutB, 16 / sizeof(ElementB), + ElementAccumulator, + MmaTileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout(sizeof(typename CollectiveEpilogue::SharedStorage))>, + MainloopSchedule + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, + CollectiveMainloop, + CollectiveEpilogue + >; + + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + auto pass = test::gemm::device::TestSmallFusion(1.0, 0.5); + EXPECT_TRUE(pass); +} + +TEST(SM100_Device_Gemm_e4m3t_e4m3n_e4m3n_tensorop_2sm_f32_bias_relu, 512x512x128_4x4x1) { + using LayoutA = cutlass::layout::RowMajor; + using LayoutB = cutlass::layout::ColumnMajor; + using LayoutC = cutlass::layout::ColumnMajor; + using ElementA = cutlass::float_e4m3_t; + using ElementB = cutlass::float_e4m3_t; + using ElementC = cutlass::float_e4m3_t; + using ElementD = cutlass::float_e4m3_t; + using ElementAccumulator = float; + using ElementCompute = float; + using ElementBias = cutlass::half_t; + using MmaTileShape = cute::Shape<_256,_128,Int<128 / sizeof(ElementA)>>; + using ClusterShape = Shape<_4,_4,_1>; + + using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecialized2Sm; + using FusionOperation = cutlass::epilogue::fusion::ScaledLinCombPerRowBiasEltAct< + cutlass::epilogue::thread::ReLU, ElementD, ElementCompute, ElementBias>; + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + MmaTileShape, ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + ElementAccumulator, ElementCompute, + ElementC, LayoutC, 16 / sizeof(ElementC), + ElementD, LayoutC, 16 / sizeof(ElementD), + EpilogueSchedule, + FusionOperation + >::CollectiveOp; + + using MainloopSchedule = cutlass::gemm::KernelTmaWarpSpecialized2SmSm100; + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + ElementA, LayoutA, 16 / sizeof(ElementA), + ElementB, LayoutB, 16 / sizeof(ElementB), + ElementAccumulator, + MmaTileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout(sizeof(typename CollectiveEpilogue::SharedStorage))>, + MainloopSchedule + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, + CollectiveMainloop, + CollectiveEpilogue + >; + + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + auto pass = test::gemm::device::TestSmallFusion(1.0, 0.5); + EXPECT_TRUE(pass); +} +#endif // #if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) diff --git a/test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_f8_f8_f8_tensor_op_f32_runtime_datatype.cu b/test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_f8_f8_f8_tensor_op_f32_runtime_datatype.cu new file mode 100644 index 00000000..b2146f53 --- /dev/null +++ b/test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_f8_f8_f8_tensor_op_f32_runtime_datatype.cu @@ -0,0 +1,295 @@ +/*************************************************************************************************** + * 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 Tests for device-wide GEMM interface +*/ + +#include + +#include "cutlass/cutlass.h" +#include "cute/tensor.hpp" +#include "cute/atom/mma_atom.hpp" + +#include "cutlass/numeric_types.h" + +#include "cutlass/gemm/device/gemm_universal_adapter.h" +#include "cutlass/gemm/kernel/gemm_universal.hpp" +#include "cutlass/gemm/collective/collective_builder.hpp" + +#include "cutlass/epilogue/dispatch_policy.hpp" +#include "cutlass/epilogue/collective/collective_builder.hpp" + +#include "cutlass/epilogue/thread/activation.h" +#include "../../../../common/cutlass_unit_test.h" + +#include "../../gemm_testbed_3x.hpp" + +using namespace cute; + +#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) + +TEST(SM100_Device_Gemm_e5m2t_e4m3n_e4m3t_tensorop_2sm_f32_runtime_datatype, 256x128x128_2x2x1) { + using CollectiveEpilogue = + typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + cute::Shape, + cute::Shape, + cutlass::epilogue::collective::EpilogueTileAuto, + float, float, + cutlass::float_e4m3_t, cutlass::layout::RowMajor, 16, + cutlass::float_e4m3_t, cutlass::layout::RowMajor, 16, + cutlass::epilogue::TmaWarpSpecialized2Sm, + + cutlass::epilogue::fusion::LinearCombination< + cutlass::float_e4m3_t, + float, + cutlass::float_e4m3_t, + float + > + + >::CollectiveOp; + + using CollectiveMainloop = + typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + cutlass::type_erased_dynamic_float8_t, cutlass::layout::RowMajor, 16, + cutlass::type_erased_dynamic_float8_t, cutlass::layout::ColumnMajor, 16, + float, + cute::Shape, + cute::Shape, + cutlass::gemm::collective::StageCountAutoCarveout, + cutlass::gemm::KernelTmaWarpSpecialized2SmSm100 + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + cute::Shape, + CollectiveMainloop, + CollectiveEpilogue, + void>; + + using namespace test::gemm::device; + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + + auto pass = TestRuntimeDataTypeSmall(cute::UMMA::MXF8F6F4Format::E5M2, cute::UMMA::MXF8F6F4Format::E4M3); + EXPECT_TRUE(pass); + +} + +TEST(SM100_Device_Gemm_e5m2t_e4m3n_e4m3t_tensorop_1sm_f32_runtime_datatype, 256x256x128_2x2x1) { + using CollectiveEpilogue = + typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + cute::Shape, + cute::Shape, + cutlass::epilogue::collective::EpilogueTileAuto, + float, float, + cutlass::float_e4m3_t, cutlass::layout::RowMajor, 16, + cutlass::float_e4m3_t, cutlass::layout::RowMajor, 16, + cutlass::epilogue::TmaWarpSpecialized1Sm, + + cutlass::epilogue::fusion::LinearCombination< + cutlass::float_e4m3_t, + float, + cutlass::float_e4m3_t, + float + > + + >::CollectiveOp; + + using CollectiveMainloop = + typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + cutlass::type_erased_dynamic_float8_t, cutlass::layout::RowMajor, 16, + cutlass::type_erased_dynamic_float8_t, cutlass::layout::ColumnMajor, 16, + float, + cute::Shape, + cute::Shape, + cutlass::gemm::collective::StageCountAutoCarveout, + cutlass::gemm::KernelTmaWarpSpecialized1SmSm100 + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + cute::Shape, + CollectiveMainloop, + CollectiveEpilogue, + void>; + + using namespace test::gemm::device; + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + + auto pass = TestRuntimeDataTypeSmall(cute::UMMA::MXF8F6F4Format::E5M2, cute::UMMA::MXF8F6F4Format::E4M3); + EXPECT_TRUE(pass); + +} + +TEST(SM100_Device_Gemm_e4m3t_e5m2n_e4m3t_tensorop_1sm_f32_runtime_datatype, 256x256x128_2x2x1) { + using CollectiveEpilogue = + typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + cute::Shape, + cute::Shape, + cutlass::epilogue::collective::EpilogueTileAuto, + float, float, + cutlass::float_e4m3_t, cutlass::layout::RowMajor, 16, + cutlass::float_e4m3_t, cutlass::layout::RowMajor, 16, + cutlass::epilogue::TmaWarpSpecialized1Sm, + + cutlass::epilogue::fusion::LinearCombination< + cutlass::float_e4m3_t, + float, + cutlass::float_e4m3_t, + float + > + + >::CollectiveOp; + + using CollectiveMainloop = + typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + cutlass::type_erased_dynamic_float8_t, cutlass::layout::RowMajor, 16, + cutlass::type_erased_dynamic_float8_t, cutlass::layout::ColumnMajor, 16, + float, + cute::Shape, + cute::Shape, + cutlass::gemm::collective::StageCountAutoCarveout, + cutlass::gemm::KernelTmaWarpSpecialized1SmSm100 + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + cute::Shape, + CollectiveMainloop, + CollectiveEpilogue, + void>; + + using namespace test::gemm::device; + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + + auto pass = TestRuntimeDataTypeSmall(cute::UMMA::MXF8F6F4Format::E4M3, cute::UMMA::MXF8F6F4Format::E5M2); + EXPECT_TRUE(pass); + +} + +TEST(SM100_Device_Gemm_e4m3t_e4m3n_e4m3t_tensorop_1sm_f32_runtime_datatype, 256x256x128_2x2x1) { + using CollectiveEpilogue = + typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + cute::Shape, + cute::Shape, + cutlass::epilogue::collective::EpilogueTileAuto, + float, float, + cutlass::float_e4m3_t, cutlass::layout::RowMajor, 16, + cutlass::float_e4m3_t, cutlass::layout::RowMajor, 16, + cutlass::epilogue::TmaWarpSpecialized1Sm, + + cutlass::epilogue::fusion::LinearCombination< + cutlass::float_e4m3_t, + float, + cutlass::float_e4m3_t, + float + > + + >::CollectiveOp; + + using CollectiveMainloop = + typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + cutlass::type_erased_dynamic_float8_t, cutlass::layout::RowMajor, 16, + cutlass::type_erased_dynamic_float8_t, cutlass::layout::ColumnMajor, 16, + float, + cute::Shape, + cute::Shape, + cutlass::gemm::collective::StageCountAutoCarveout, + cutlass::gemm::KernelTmaWarpSpecialized1SmSm100 + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + cute::Shape, + CollectiveMainloop, + CollectiveEpilogue, + void>; + + using namespace test::gemm::device; + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + + auto pass = TestRuntimeDataTypeSmall(cute::UMMA::MXF8F6F4Format::E4M3, cute::UMMA::MXF8F6F4Format::E4M3); + EXPECT_TRUE(pass); + +} + +TEST(SM100_Device_Gemm_e5m2t_e5m2n_e5m2t_tensorop_2sm_f32_runtime_datatype, 256x256x128_2x2x1) { + using CollectiveEpilogue = + typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + cute::Shape, + cute::Shape, + cutlass::epilogue::collective::EpilogueTileAuto, + float, float, + cutlass::float_e5m2_t, cutlass::layout::RowMajor, 16, + cutlass::float_e5m2_t, cutlass::layout::RowMajor, 16, + cutlass::epilogue::TmaWarpSpecialized1Sm, + + cutlass::epilogue::fusion::LinearCombination< + cutlass::float_e5m2_t, + float, + cutlass::float_e5m2_t, + float + > + + >::CollectiveOp; + + using CollectiveMainloop = + typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + cutlass::type_erased_dynamic_float8_t, cutlass::layout::RowMajor, 16, + cutlass::type_erased_dynamic_float8_t, cutlass::layout::ColumnMajor, 16, + float, + cute::Shape, + cute::Shape, + cutlass::gemm::collective::StageCountAutoCarveout, + cutlass::gemm::KernelTmaWarpSpecialized2SmSm100 + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + cute::Shape, + CollectiveMainloop, + CollectiveEpilogue, + void>; + + using namespace test::gemm::device; + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + + auto pass = TestRuntimeDataTypeSmall(cute::UMMA::MXF8F6F4Format::E5M2, cute::UMMA::MXF8F6F4Format::E5M2); + EXPECT_TRUE(pass); + +} + +#endif // #if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) diff --git a/test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_f8_f8_f8_tensor_op_s32_batch_alpha_beta.cu b/test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_f8_f8_f8_tensor_op_s32_batch_alpha_beta.cu new file mode 100644 index 00000000..bb90bf87 --- /dev/null +++ b/test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_f8_f8_f8_tensor_op_s32_batch_alpha_beta.cu @@ -0,0 +1,219 @@ +/*************************************************************************************************** + * 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 Tests for device-wide GEMM interface +*/ + +#include + +#include "cutlass/cutlass.h" +#include "cute/tensor.hpp" +#include "cute/atom/mma_atom.hpp" + +#include "cutlass/numeric_types.h" + +#include "cutlass/gemm/device/gemm_universal_adapter.h" +#include "cutlass/gemm/kernel/gemm_universal.hpp" +#include "cutlass/gemm/collective/collective_builder.hpp" +#include "cutlass/epilogue/dispatch_policy.hpp" +#include "cutlass/epilogue/collective/collective_builder.hpp" +#include "cutlass/epilogue/thread/linear_combination.h" + +#include "../../../../common/cutlass_unit_test.h" + +#include "../../gemm_testbed_3x.hpp" + +using namespace cute; + +#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) + +/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// +////////////////////////////////////////// Test Batch alpha and beta ////////////////////////////////////////// +/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// + +TEST(SM100_Device_Gemm_e4m3t_e4m3n_e4m3n_tensorop_1cta_s32_batch_alpha_beta, 128x64x128_1x1x1) { + using LayoutA = cutlass::layout::RowMajor; + using LayoutB = cutlass::layout::ColumnMajor; + using LayoutC = cutlass::layout::ColumnMajor; + using ElementA = cutlass::float_e4m3_t; + using ElementB = cutlass::float_e4m3_t; + using ElementC = cutlass::float_e4m3_t; + using ElementD = cutlass::float_e4m3_t; + using ElementAccumulator = float; + using ElementCompute = float; + using ElementBias = cutlass::half_t; + using MmaTileShape = cute::Shape<_128,_64,Int<128 / sizeof(ElementA)>>; + using ClusterShape = Shape<_1,_1,_1>; + + using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecialized1Sm; + + using FusionOperation = cutlass::epilogue::fusion::LinearCombination< + ElementD, + ElementCompute, + ElementC, + ElementBias + >; + + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + MmaTileShape, ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + ElementAccumulator, ElementCompute, + ElementC, LayoutC, 16 / sizeof(ElementC), + ElementD, LayoutC, 16 / sizeof(ElementD), + EpilogueSchedule, + FusionOperation + >::CollectiveOp; + + using MainloopSchedule = cutlass::gemm::KernelTmaWarpSpecialized1SmSm100; + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + ElementA, LayoutA, 16 / sizeof(ElementA), + ElementB, LayoutB, 16 / sizeof(ElementB), + ElementAccumulator, + MmaTileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout(sizeof(typename CollectiveEpilogue::SharedStorage))>, + MainloopSchedule + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, + CollectiveMainloop, + CollectiveEpilogue + >; + + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + auto pass = test::gemm::device::TestSmallFusion(1.0, 1.0); // beta is [1.0, 2.0] + EXPECT_TRUE(pass); +} + +TEST(SM100_Device_Gemm_e4m3t_e4m3n_e4m3n_tensorop_1sm_f32_bias_relu_batch_alpha_beta, 128x128x128_1x1x1) { + using LayoutA = cutlass::layout::RowMajor; + using LayoutB = cutlass::layout::ColumnMajor; + using LayoutC = cutlass::layout::ColumnMajor; + using ElementA = cutlass::float_e4m3_t; + using ElementB = cutlass::float_e4m3_t; + using ElementC = cutlass::float_e4m3_t; + using ElementD = cutlass::float_e4m3_t; + using ElementAccumulator = float; + using ElementCompute = float; + using ElementBias = cutlass::half_t; + using MmaTileShape = cute::Shape<_128,_128,Int<128 / sizeof(ElementA)>>; + using ClusterShape = Shape<_1,_1,_1>; + + using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecialized1Sm; + using FusionOperation = cutlass::epilogue::fusion::ScaledLinCombPerRowBiasEltAct< + cutlass::epilogue::thread::ReLU, ElementD, ElementCompute, ElementBias>; + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + MmaTileShape, ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + ElementAccumulator, ElementCompute, + ElementC, LayoutC, 16 / sizeof(ElementC), + ElementD, LayoutC, 16 / sizeof(ElementD), + EpilogueSchedule, + FusionOperation + >::CollectiveOp; + + using MainloopSchedule = cutlass::gemm::KernelTmaWarpSpecialized1SmSm100; + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + ElementA, LayoutA, 16 / sizeof(ElementA), + ElementB, LayoutB, 16 / sizeof(ElementB), + ElementAccumulator, + MmaTileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout(sizeof(typename CollectiveEpilogue::SharedStorage))>, + MainloopSchedule + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, + CollectiveMainloop, + CollectiveEpilogue + >; + + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + auto pass = test::gemm::device::TestSmallFusion(1.0, 0.5); // beta is [0.5, 1.5] + EXPECT_TRUE(pass); +} + +TEST(SM100_Device_Gemm_e4m3t_e4m3n_e4m3n_tensorop_1sm_f32_bias_relu__batch_alpha_beta0, 128x128x128_1x1x1) { + using LayoutA = cutlass::layout::RowMajor; + using LayoutB = cutlass::layout::ColumnMajor; + using LayoutC = cutlass::layout::ColumnMajor; + using ElementA = cutlass::float_e4m3_t; + using ElementB = cutlass::float_e4m3_t; + using ElementC = cutlass::float_e4m3_t; + using ElementD = cutlass::float_e4m3_t; + using ElementAccumulator = float; + using ElementCompute = float; + using ElementBias = cutlass::half_t; + using MmaTileShape = cute::Shape<_128,_128,Int<128 / sizeof(ElementA)>>; + using ClusterShape = Shape<_1,_1,_1>; + + using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecialized1Sm; + using FusionOperation = cutlass::epilogue::fusion::ScaledLinCombPerRowBiasEltAct< + cutlass::epilogue::thread::ReLU, ElementD, ElementCompute, ElementBias>; + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + MmaTileShape, ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + ElementAccumulator, ElementCompute, + ElementC, LayoutC, 16 / sizeof(ElementC), + ElementD, LayoutC, 16 / sizeof(ElementD), + EpilogueSchedule, + FusionOperation + >::CollectiveOp; + + using MainloopSchedule = cutlass::gemm::KernelTmaWarpSpecialized1SmSm100; + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + ElementA, LayoutA, 16 / sizeof(ElementA), + ElementB, LayoutB, 16 / sizeof(ElementB), + ElementAccumulator, + MmaTileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout(sizeof(typename CollectiveEpilogue::SharedStorage))>, + MainloopSchedule + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, + CollectiveMainloop, + CollectiveEpilogue + >; + + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + auto pass = test::gemm::device::TestSmallFusion(1.0, -1.0); // beta is [-1.0, 0.0] + EXPECT_TRUE(pass); +} + +#endif // #if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) diff --git a/test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_i8_i8_i8_tensor_op_s32_bias_relu.cu b/test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_i8_i8_i8_tensor_op_s32_bias_relu.cu new file mode 100644 index 00000000..c5a1c46d --- /dev/null +++ b/test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_i8_i8_i8_tensor_op_s32_bias_relu.cu @@ -0,0 +1,278 @@ +/*************************************************************************************************** + * 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 Tests for device-wide GEMM interface +*/ + +#include + +#include "cutlass/cutlass.h" +#include "cute/tensor.hpp" +#include "cute/atom/mma_atom.hpp" + +#include "cutlass/numeric_types.h" +#include "cutlass/arch/mma_sm100.h" + +#include "cutlass/gemm/device/gemm_universal_adapter.h" +#include "cutlass/gemm/kernel/gemm_universal.hpp" +#include "cutlass/gemm/collective/collective_builder.hpp" +#include "cutlass/epilogue/dispatch_policy.hpp" +#include "cutlass/epilogue/collective/collective_builder.hpp" +#include "../../../../common/cutlass_unit_test.h" + +#include "../../gemm_testbed_3x.hpp" + +using namespace cute; + +#if (defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) && !defined(CUTLASS_SM100_FAMILY_ARCHS_ENABLED)) + +/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// +///////////////////////////////////////////// 128x64x128 1x1x1 TMEM 4x1 //////////////////////////////////////////// +/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// + +TEST(SM100_Device_Gemm_s8t_s8n_s8n_tensorop_1cta_s32_bias_relu, 128x64x128_1x1x1) { + using LayoutA = cutlass::layout::RowMajor; + using LayoutB = cutlass::layout::ColumnMajor; + using LayoutC = cutlass::layout::ColumnMajor; + using ElementA = int8_t; + using ElementB = int8_t; + using ElementC = int8_t; + using ElementD = int8_t; + using ElementAccumulator = int32_t; + using ElementCompute = float; + using ElementBias = int8_t; + using MmaTileShape = cute::Shape<_128,_64,Int<128 / sizeof(ElementA)>>; + using ClusterShape = Shape<_1,_1,_1>; + + using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecialized1Sm; + using FusionOperation = cutlass::epilogue::fusion::LinCombPerRowBiasEltAct< + cutlass::epilogue::thread::ReLu, ElementD, ElementCompute, ElementBias>; + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + MmaTileShape, ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + ElementAccumulator, ElementCompute, + ElementC, LayoutC, 16 / sizeof(ElementC), + ElementD, LayoutC, 16 / sizeof(ElementD), + EpilogueSchedule, + FusionOperation + >::CollectiveOp; + + using MainloopSchedule = cutlass::gemm::KernelTmaWarpSpecialized1SmSm100; + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + ElementA, LayoutA, 16 / sizeof(ElementA), + ElementB, LayoutB, 16 / sizeof(ElementB), + ElementAccumulator, + MmaTileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout(sizeof(typename CollectiveEpilogue::SharedStorage))>, + MainloopSchedule + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, + CollectiveMainloop, + CollectiveEpilogue + >; + + using namespace test::gemm::device; + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + auto pass = TestSmallFusion(2, 0.5, CheckEquality::EXACT); + EXPECT_TRUE(pass); +} +/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// +///////////////////////////////////////////// 128x64x128 4x2x1 TMEM 4x1 //////////////////////////////////////////// +/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// + +TEST(SM100_Device_Gemm_s8t_s8n_s8n_tensorop_1cta_s32_bias_relu, 512x128x128_4x2x1) { + using LayoutA = cutlass::layout::RowMajor; + using LayoutB = cutlass::layout::ColumnMajor; + using LayoutC = cutlass::layout::ColumnMajor; + using ElementA = int8_t; + using ElementB = int8_t; + using ElementC = int8_t; + using ElementD = int8_t; + using ElementAccumulator = int32_t; + using ElementCompute = float; + using ElementBias = int8_t; + using MmaTileShape = Shape<_128,_64,Int<128 / sizeof(ElementA)>>; + using ClusterShape = Shape<_4,_2,_1>; + + using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecialized1Sm; + using FusionOperation = cutlass::epilogue::fusion::LinCombPerRowBiasEltAct< + cutlass::epilogue::thread::ReLu, ElementD, ElementCompute, ElementBias>; + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + MmaTileShape, ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + ElementAccumulator, ElementCompute, + ElementC, LayoutC, 16 / sizeof(ElementC), + ElementD, LayoutC, 16 / sizeof(ElementD), + EpilogueSchedule, + FusionOperation + >::CollectiveOp; + + using MainloopSchedule = cutlass::gemm::KernelTmaWarpSpecialized1SmSm100; + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + ElementA, LayoutA, 16 / sizeof(ElementA), + ElementB, LayoutB, 16 / sizeof(ElementB), + ElementAccumulator, + MmaTileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout(sizeof(typename CollectiveEpilogue::SharedStorage))>, + MainloopSchedule + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, + CollectiveMainloop, + CollectiveEpilogue + >; + + using namespace test::gemm::device; + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + auto pass = TestSmallFusion(2, 0.5, CheckEquality::EXACT); + EXPECT_TRUE(pass); +} + +/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// +///////////////////////////////////////////// 64x256x128 1x1x1 TMEM 4x1 //////////////////////////////////////////// +/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// + +TEST(SM100_Device_Gemm_s8t_s8n_s32n_tensorop_1cta_s32_bias_relu, 64x256x128_1x1x1) { + using LayoutA = cutlass::layout::RowMajor; + using LayoutB = cutlass::layout::ColumnMajor; + using LayoutC = cutlass::layout::ColumnMajor; + using ElementA = int8_t; + using ElementB = int8_t; + using ElementC = int32_t; + using ElementD = int32_t; + using ElementAccumulator = int32_t; + using ElementCompute = int32_t; + using ElementBias = int32_t; + using MmaTileShape = cute::Shape<_64,_256,Int<128 / sizeof(ElementA)>>; + using ClusterShape = Shape<_1,_1,_1>; + + using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecialized1Sm; + using FusionOperation = cutlass::epilogue::fusion::LinCombPerRowBiasEltAct< + cutlass::epilogue::thread::ReLu, ElementD, ElementCompute, ElementBias>; + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + MmaTileShape, ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + ElementAccumulator, ElementCompute, + ElementC, LayoutC, 16 / sizeof(ElementC), + ElementD, LayoutC, 16 / sizeof(ElementD), + EpilogueSchedule, + FusionOperation + >::CollectiveOp; + + using MainloopSchedule = cutlass::gemm::KernelTmaWarpSpecialized1SmSm100; + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + ElementA, LayoutA, 16 / sizeof(ElementA), + ElementB, LayoutB, 16 / sizeof(ElementB), + ElementAccumulator, + MmaTileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout(sizeof(typename CollectiveEpilogue::SharedStorage))>, + MainloopSchedule + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, + CollectiveMainloop, + CollectiveEpilogue + >; + + using namespace test::gemm::device; + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + auto pass = TestSmallFusion(2, 0.5, CheckEquality::EXACT); + EXPECT_TRUE(pass); +} + +/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// +///////////////////////////////////////////// 64x256x128 2x4x1 TMEM 2x2 //////////////////////////////////////////// +/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// + +TEST(SM100_Device_Gemm_s8t_s8n_s8n_tensorop_2cta_s32_bias_relu, 128x1024x128_2x4x1) { + using LayoutA = cutlass::layout::RowMajor; + using LayoutB = cutlass::layout::ColumnMajor; + using LayoutC = cutlass::layout::ColumnMajor; + using ElementA = int8_t; + using ElementB = int8_t; + using ElementC = int8_t; + using ElementD = int8_t; + using ElementAccumulator = int32_t; + using ElementCompute = float; + using ElementBias = int8_t; + using MmaTileShape = Shape<_128,_256,Int<128 / sizeof(ElementA)>>; + using ClusterShape = Shape<_2,_4,_1>; + + using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecialized2Sm; + using FusionOperation = cutlass::epilogue::fusion::LinCombPerRowBiasEltAct< + cutlass::epilogue::thread::ReLu, ElementD, ElementCompute, ElementBias>; + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + MmaTileShape, ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + ElementAccumulator, ElementCompute, + ElementC, LayoutC, 16 / sizeof(ElementC), + ElementD, LayoutC, 16 / sizeof(ElementD), + EpilogueSchedule, + FusionOperation + >::CollectiveOp; + + using MainloopSchedule = cutlass::gemm::KernelTmaWarpSpecialized2SmSm100; + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + ElementA, LayoutA, 16 / sizeof(ElementA), + ElementB, LayoutB, 16 / sizeof(ElementB), + ElementAccumulator, + MmaTileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout(sizeof(typename CollectiveEpilogue::SharedStorage))>, + MainloopSchedule + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, + CollectiveMainloop, + CollectiveEpilogue + >; + + using namespace test::gemm::device; + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + auto pass = TestSmallFusion(2, 0.5, CheckEquality::EXACT); + EXPECT_TRUE(pass); +} + +#endif // #if (defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) && !defined(CUTLASS_SM100_FAMILY_ARCHS_ENABLED)) + diff --git a/test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_i8_i8_i8_tensor_op_s32_vector_alpha_beta.cu b/test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_i8_i8_i8_tensor_op_s32_vector_alpha_beta.cu new file mode 100644 index 00000000..7ac363d0 --- /dev/null +++ b/test/unit/gemm/device/sm100_tensorop_gemm/extra_tests/sm100_gemm_i8_i8_i8_tensor_op_s32_vector_alpha_beta.cu @@ -0,0 +1,343 @@ +/*************************************************************************************************** + * 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 Tests for device-wide GEMM interface +*/ + +#include + +#include "cutlass/cutlass.h" +#include "cute/tensor.hpp" +#include "cute/atom/mma_atom.hpp" + +#include "cutlass/numeric_types.h" +#include "cutlass/arch/mma_sm100.h" + +#include "cutlass/gemm/device/gemm_universal_adapter.h" +#include "cutlass/gemm/kernel/gemm_universal.hpp" +#include "cutlass/gemm/collective/collective_builder.hpp" +#include "cutlass/epilogue/dispatch_policy.hpp" +#include "cutlass/epilogue/collective/collective_builder.hpp" +#include "cutlass/epilogue/thread/linear_combination.h" + +#include "../../../../common/cutlass_unit_test.h" + +#include "../../gemm_testbed_3x.hpp" + +using namespace cute; + +#if (defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) && !defined(CUTLASS_SM100_FAMILY_ARCHS_ENABLED)) + +/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// +////////////////////////////////////////// Test Vector alpha and vector beta ////////////////////////////////////////// +/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// + +/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// +///////////////////////////////////////////// 128x64x128 1x1x1 TMEM 4x1 //////////////////////////////////////////// +/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// + +TEST(SM100_Device_Gemm_s8t_s8n_s8n_tensorop_1cta_s32_vector_alpha_beta, 128x64x128_1x1x1) { + using LayoutA = cutlass::layout::RowMajor; + using LayoutB = cutlass::layout::ColumnMajor; + using LayoutC = cutlass::layout::ColumnMajor; + using ElementA = int8_t; + using ElementB = int8_t; + using ElementC = int8_t; + using ElementD = int8_t; + using ElementAccumulator = int32_t; + using ElementCompute = float; + using ElementBias = int8_t; + using MmaTileShape = cute::Shape<_128,_64,Int<128 / sizeof(ElementA)>>; + using ClusterShape = Shape<_1,_1,_1>; + + using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecialized1Sm; + using FusionOperation = cutlass::epilogue::fusion::PerRowLinCombPerRowBiasEltAct< + cutlass::epilogue::thread::ReLU, ElementD, ElementCompute, ElementBias>; + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + MmaTileShape, ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + ElementAccumulator, ElementCompute, + ElementC, LayoutC, 16 / sizeof(ElementC), + ElementD, LayoutC, 16 / sizeof(ElementD), + EpilogueSchedule, + FusionOperation + >::CollectiveOp; + + using MainloopSchedule = cutlass::gemm::KernelTmaWarpSpecialized1SmSm100; + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + ElementA, LayoutA, 16 / sizeof(ElementA), + ElementB, LayoutB, 16 / sizeof(ElementB), + ElementAccumulator, + MmaTileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout(sizeof(typename CollectiveEpilogue::SharedStorage))>, + MainloopSchedule + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, + CollectiveMainloop, + CollectiveEpilogue + >; + + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + auto pass = test::gemm::device::TestSmallFusion(1.0, 1.0); + EXPECT_TRUE(pass); +} + +/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// +///////////////////////////////////////////// 128x64x128 4x2x1 TMEM 4x1 //////////////////////////////////////////// +/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// + +TEST(SM100_Device_Gemm_s8t_s8n_s8n_tensorop_1cta_s32_vector_alpha_beta, 512x128x128_4x2x1) { + using LayoutA = cutlass::layout::RowMajor; + using LayoutB = cutlass::layout::ColumnMajor; + using LayoutC = cutlass::layout::ColumnMajor; + using ElementA = int8_t; + using ElementB = int8_t; + using ElementC = int8_t; + using ElementD = int8_t; + using ElementAccumulator = int32_t; + using ElementCompute = float; + using ElementBias = int8_t; + using MmaTileShape = cute::Shape<_128,_64,Int<128 / sizeof(ElementA)>>; + using ClusterShape = Shape<_4,_2,_1>; + + using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecialized1Sm; + using FusionOperation = cutlass::epilogue::fusion::PerRowLinCombPerRowBiasEltAct< + cutlass::epilogue::thread::ReLU, ElementD, ElementCompute, ElementBias>; + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + MmaTileShape, ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + ElementAccumulator, ElementCompute, + ElementC, LayoutC, 16 / sizeof(ElementC), + ElementD, LayoutC, 16 / sizeof(ElementD), + EpilogueSchedule, + FusionOperation + >::CollectiveOp; + + using MainloopSchedule = cutlass::gemm::KernelTmaWarpSpecialized1SmSm100; + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + ElementA, LayoutA, 16 / sizeof(ElementA), + ElementB, LayoutB, 16 / sizeof(ElementB), + ElementAccumulator, + MmaTileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout(sizeof(typename CollectiveEpilogue::SharedStorage))>, + MainloopSchedule + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, + CollectiveMainloop, + CollectiveEpilogue + >; + + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + auto pass = test::gemm::device::TestSmallFusion(1.0, 1.0); + EXPECT_TRUE(pass); +} + +/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// +///////////////////////////////////////////// 64x256x128 1x1x1 TMEM 4x1 //////////////////////////////////////////// +/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// + +TEST(SM100_Device_Gemm_s8t_s8n_s8n_tensorop_1cta_s32_vector_alpha_beta, 64x256x128_1x1x1) { + using LayoutA = cutlass::layout::RowMajor; + using LayoutB = cutlass::layout::ColumnMajor; + using LayoutC = cutlass::layout::ColumnMajor; + using ElementA = int8_t; + using ElementB = int8_t; + using ElementC = int8_t; + using ElementD = int8_t; + using ElementAccumulator = int32_t; + using ElementCompute = float; + using ElementBias = int8_t; + using MmaTileShape = cute::Shape<_64,_256,Int<128 / sizeof(ElementA)>>; + using ClusterShape = Shape<_1,_1,_1>; + + using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecialized1Sm; + using FusionOperation = cutlass::epilogue::fusion::PerRowLinCombPerRowBiasEltAct< + cutlass::epilogue::thread::ReLU, ElementD, ElementCompute, ElementBias>; + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + MmaTileShape, ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + ElementAccumulator, ElementCompute, + ElementC, LayoutC, 16 / sizeof(ElementC), + ElementD, LayoutC, 16 / sizeof(ElementD), + EpilogueSchedule, + FusionOperation + >::CollectiveOp; + + using MainloopSchedule = cutlass::gemm::KernelTmaWarpSpecialized1SmSm100; + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + ElementA, LayoutA, 16 / sizeof(ElementA), + ElementB, LayoutB, 16 / sizeof(ElementB), + ElementAccumulator, + MmaTileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout(sizeof(typename CollectiveEpilogue::SharedStorage))>, + MainloopSchedule + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, + CollectiveMainloop, + CollectiveEpilogue + >; + + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + auto pass = test::gemm::device::TestSmallFusion(1.0, 1.0); + EXPECT_TRUE(pass); +} + +/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// +///////////////////////////////////////////// 64x256x128 2x4x1 TMEM 2x2 //////////////////////////////////////////// +/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// + +TEST(SM100_Device_Gemm_s8t_s8n_s8n_tensorop_1cta_s32_vector_alpha_beta, 128x1024x128_2x4x1) { + using LayoutA = cutlass::layout::RowMajor; + using LayoutB = cutlass::layout::ColumnMajor; + using LayoutC = cutlass::layout::ColumnMajor; + using ElementA = int8_t; + using ElementB = int8_t; + using ElementC = int8_t; + using ElementD = int8_t; + using ElementAccumulator = int32_t; + using ElementCompute = float; + using ElementBias = int8_t; + using MmaTileShape = cute::Shape<_64,_256,Int<128 / sizeof(ElementA)>>; + using ClusterShape = Shape<_2,_4,_1>; + + using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecialized1Sm; + using FusionOperation = cutlass::epilogue::fusion::PerRowLinCombPerRowBiasEltAct< + cutlass::epilogue::thread::ReLU, ElementD, ElementCompute, ElementBias>; + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + MmaTileShape, ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + ElementAccumulator, ElementCompute, + ElementC, LayoutC, 16 / sizeof(ElementC), + ElementD, LayoutC, 16 / sizeof(ElementD), + EpilogueSchedule, + FusionOperation + >::CollectiveOp; + + using MainloopSchedule = cutlass::gemm::KernelTmaWarpSpecialized1SmSm100; + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + ElementA, LayoutA, 16 / sizeof(ElementA), + ElementB, LayoutB, 16 / sizeof(ElementB), + ElementAccumulator, + MmaTileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout(sizeof(typename CollectiveEpilogue::SharedStorage))>, + MainloopSchedule + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, + CollectiveMainloop, + CollectiveEpilogue + >; + + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + auto pass = test::gemm::device::TestSmallFusion(1.0, 1.0); + EXPECT_TRUE(pass); +} + +/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// +//////////////////////////////// Dynamic vector/scalar broadcast /////////////////////////// +/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// + +TEST(SM100_Device_Gemm_s8t_s8n_s8n_tensorop_1cta_s32_dynamic_vector_alpha_beta, 128x64x128_1x1x1) { + using LayoutA = cutlass::layout::RowMajor; + using LayoutB = cutlass::layout::ColumnMajor; + using LayoutC = cutlass::layout::ColumnMajor; + using ElementA = int8_t; + using ElementB = int8_t; + using ElementC = int8_t; + using ElementD = int8_t; + using ElementAccumulator = int32_t; + using ElementCompute = float; + using ElementBias = int8_t; + using MmaTileShape = cute::Shape<_128,_64,Int<128 / sizeof(ElementA)>>; + using ClusterShape = Shape<_1,_1,_1>; + + using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecialized1Sm; + using FusionOperation = cutlass::epilogue::fusion::PerRowLinCombPerRowBiasEltAct< + cutlass::epilogue::thread::ReLU, ElementD, ElementCompute, ElementBias>; + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + MmaTileShape, ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + ElementAccumulator, ElementCompute, + ElementC, LayoutC, 16 / sizeof(ElementC), + ElementD, LayoutC, 16 / sizeof(ElementD), + EpilogueSchedule, + FusionOperation + >::CollectiveOp; + + using MainloopSchedule = cutlass::gemm::KernelTmaWarpSpecialized1SmSm100; + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + ElementA, LayoutA, 16 / sizeof(ElementA), + ElementB, LayoutB, 16 / sizeof(ElementB), + ElementAccumulator, + MmaTileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout(sizeof(typename CollectiveEpilogue::SharedStorage))>, + MainloopSchedule + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, + CollectiveMainloop, + CollectiveEpilogue + >; + + using namespace test::gemm::device; + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + constexpr bool force_legacy_epilogue = false; + constexpr bool apply_alignment_offset = false; + // non-batched host scalar, beta 0 + EXPECT_TRUE((TestSmallFusion(1.0, 0.0, CheckEquality::EXACT, ScalarLoc::ON_HOST, VectorScale::DISABLED))); + // non-batched host scalar, beta 1 + EXPECT_TRUE((TestSmallFusion(1.0, 1.0, CheckEquality::EXACT, ScalarLoc::ON_HOST, VectorScale::DISABLED))); + // batched device scalar, beta 0 + EXPECT_TRUE((TestSmallFusion(1.0, 0.0, CheckEquality::EXACT, ScalarLoc::ON_DEVICE, VectorScale::DISABLED))); + // batched device scalar, beta 1 + EXPECT_TRUE((TestSmallFusion(1.0, 1.0, CheckEquality::EXACT, ScalarLoc::ON_DEVICE, VectorScale::DISABLED))); +} + +#endif // #if (defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) && !defined(CUTLASS_SM100_FAMILY_ARCHS_ENABLED))