# 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. 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 import cutlass.pipeline as pipeline import cutlass.utils as utils """ A HSTU attention forward pass example for NVIDIA Ampere SM80 architecture using Cute DSL, based on the example of flash_attention_v2 for Ampere. The example showcases an implementation of HSTU attention(https://arxiv.org/abs/2402.17152) within generative recommender system. It utilize the formula: `mask(silu(q@k+rab))@v`. The implementation includes the following features: - efficient fast sigmoid implementation - block rasterization to improve L2 cache hit rate. - The correct approach to verify the results of the HSTU attention with a Pytorch implementation. To run this example: .. code-block:: bash python examples/ampere/hstu_attention.py --batch_size 4 --seqlen_q 8192 --seqlen_kv 8192 --num_head 4 --head_dim 128 --m_block_size 128 --n_block_size 64 --is_causal --perf_test The above example tests the performance of HSTU attention with batch size 4, sequence length 8192, 4 attention heads, and head dimension 128. The m_block_size is 128, and n_block_size is 64. The causal masking is enabled. There are some constraints for this example: * Only Float16 and BFloat16 are supported. * The contiguous dimension of each tensor must be at least 16 bytes aligned. * The values of `m_block_size`, `n_block_size`, and `head_dim` must be selected to stay within shared memory capacity limits. * `m_block_size * 2` must be divisible by `num_threads`, otherwise the kernel will not be able to get the correct result. * "seqlen_kv should be greater or equal to seqlen_q. """ class HSTUAttentionForwardAmpere(object): def __init__( self, dtype, batch_size, seqlen_q, seqlen_kv, num_head, head_dim, m_block_size=128, n_block_size=128, num_threads=128, enable_fast_sigmoid=False, enable_block_rasterization=False, is_causal=False, ): self._dtype = dtype self._batch_size = batch_size self._seqlen_q = seqlen_q self._seqlen_kv = seqlen_kv self._num_head = num_head self._head_dim = head_dim self._m_block_size = m_block_size self._n_block_size = n_block_size # padded head_dim to 32 for cta tile. self._head_dim_padded = (head_dim + 31) // 32 * 32 self._num_threads = num_threads self._enable_fast_sigmoid = enable_fast_sigmoid self._enable_block_rasterization = enable_block_rasterization self._is_causal = is_causal assert self._dtype == cutlass.Float16 or self._dtype == cutlass.BFloat16, ( "Only Float16 or BFloat16 is supported" ) assert self._head_dim % 8 == 0, "head dim should be multiply of 8" assert self._num_threads % 32 == 0, "num_threads should be multiply of 32" assert self._m_block_size * self._head_dim_padded // self._num_threads >= 8, ( "Small m_block_size and too many threads" ) assert self._n_block_size * self._head_dim_padded // self._num_threads >= 8, ( "Small n_block_size and too many threads" ) assert seqlen_kv >= seqlen_q, "seqlen_kv should be greater or equal to seqlen_q" self.cta_sync_barrier = pipeline.NamedBarrier( barrier_id=1, num_threads=num_threads ) @cute.jit def __call__( self, mQ: cute.Tensor, mK: cute.Tensor, mV: cute.Tensor, mO: cute.Tensor, mRAB: cute.Tensor, stream: cuda.CUstream, ): """Configures and launches the HSTU attention kernel. mQ/mK/mV/mO/mRAB has same data types(supports fp16 and bf16). mQ has layout: (batch_size, seqlen_q, num_head, head_dim):(seqlen_q * num_head * head_dim, num_head * head_dim, head_dim, 1) mK/mV/mO has same layout: (batch_size, seqlen_kv, num_head, head_dim):(seqlen_kv * num_head * head_dim, num_head * head_dim, head_dim, 1) mRAB has layout: (batch_size, num_head, seqlen_q, seqlen_kv):(seqlen_q*seqlen_kv*num_head, seqlen_q*seqlen_kv, seqlen_kv, 1) Prepares the shared memory layout, tiled copy atoms, tiled mma and shared memory storage. Then launches the kernel function with the prepared parameters. :param mQ: query tensor :type mQ: cute.Tensor :param mK: key tensor :type mK: cute.Tensor :param mV: value tensor :type mV: cute.Tensor :param mO: output tensor :type mO: cute.Tensor :param mRAB: RAB tensor :type mRAB: cute.Tensor """ # /////////////////////////////////////////////////////////////////////////////// # Shared memory layout: Q/K/V/RAB # /////////////////////////////////////////////////////////////////////////////// smem_k_block_size = 64 if self._head_dim_padded % 64 == 0 else 32 swizzle_bits = 3 if smem_k_block_size == 64 else 2 sQ_layout_atom = cute.make_composed_layout( cute.make_swizzle(swizzle_bits, 4, 3), 0, cute.make_layout((8, smem_k_block_size), stride=(smem_k_block_size, 1)), ) sQ_layout = cute.tile_to_shape( sQ_layout_atom, (self._m_block_size, self._head_dim_padded), (0, 1), ) sKV_layout_atom = sQ_layout_atom sKV_layout = cute.tile_to_shape( sKV_layout_atom, (self._n_block_size, self._head_dim_padded), (0, 1), ) sRAB_layout_atom = sQ_layout_atom sRAB_layout = cute.tile_to_shape( sRAB_layout_atom, (self._m_block_size, self._n_block_size), (0, 1) ) sO_layout = sQ_layout @cute.struct class SharedStorage: sQ: cute.struct.Align[ cute.struct.MemRange[self._dtype, cute.cosize(sQ_layout)], 1024 ] sK: cute.struct.Align[ cute.struct.MemRange[self._dtype, cute.cosize(sKV_layout)], 1024 ] sV: cute.struct.Align[ cute.struct.MemRange[self._dtype, cute.cosize(sKV_layout)], 1024 ] sRAB: cute.struct.Align[ cute.struct.MemRange[self._dtype, cute.cosize(sRAB_layout)], 1024 ] assert SharedStorage.size_in_bytes() < utils.get_smem_capacity_in_bytes( "sm_80" ), "insufficient shared memory" # /////////////////////////////////////////////////////////////////////////////// # GMEM Tiled copy: # /////////////////////////////////////////////////////////////////////////////// # Thread layouts for copies universal_copy_bits = 128 async_copy_elems = universal_copy_bits // self._dtype.width # atom_async_copy: async copy atom for QKV load atom_async_copy = cute.make_copy_atom( cute.nvgpu.cpasync.CopyG2SOp(), self._dtype, num_bits_per_copy=universal_copy_bits, ) # atom_universal_copy: universal copy atom for O store atom_universal_copy = cute.make_copy_atom( cute.nvgpu.CopyUniversalOp(), self._dtype, num_bits_per_copy=universal_copy_bits, ) # tQKV_layout: thread layout for QKV load tQKV_shape_dim_1 = sQ_layout_atom.outer.shape[1] // async_copy_elems tQKV_layout = cute.make_layout( (self._num_threads // tQKV_shape_dim_1, tQKV_shape_dim_1), stride=(tQKV_shape_dim_1, 1), ) # tO_layout: thread layout for O store tO_layout = tQKV_layout # Value layouts for copies vQKV_layout = cute.make_layout((1, async_copy_elems)) vO_layout = vQKV_layout # gmem_tiled_copy_QKV: tiled copy for QKV load gmem_tiled_copy_QKV = cute.make_tiled_copy_tv( atom_async_copy, tQKV_layout, vQKV_layout ) # gmem_tiled_copy_O: tiled copy for O store gmem_tiled_copy_O = cute.make_tiled_copy_tv( atom_universal_copy, tO_layout, vO_layout ) # /////////////////////////////////////////////////////////////////////////////// # Tiled mma # /////////////////////////////////////////////////////////////////////////////// tiled_mma = cute.make_tiled_mma( cute.nvgpu.warp.MmaF16BF16Op(self._dtype, cutlass.Float32, (16, 8, 16)), (self._num_threads // 32, 1, 1), permutation_mnk=(self._num_threads // 32 * 16, 16, 16), ) # block rasterization if cutlass.const_expr(self._enable_block_rasterization): grid_dim = ( self._batch_size, self._num_head, cute.ceil_div(mQ.shape[1], self._m_block_size), ) else: grid_dim = ( cute.ceil_div(mQ.shape[1], self._m_block_size), self._batch_size, self._num_head, ) self.kernel( mQ, mK, mV, mO, mRAB, sQ_layout, sKV_layout, sRAB_layout, sO_layout, gmem_tiled_copy_QKV, gmem_tiled_copy_O, tiled_mma, SharedStorage, ).launch( grid=grid_dim, block=[self._num_threads, 1, 1], smem=SharedStorage.size_in_bytes(), stream=stream, ) @cute.kernel def kernel( self, mQ: cute.Tensor, mK: cute.Tensor, mV: cute.Tensor, mO: cute.Tensor, mRAB: cute.Tensor, sQ_layout: cute.ComposedLayout, sKV_layout: cute.ComposedLayout, sRAB_layout: cute.ComposedLayout, sO_layout: cute.ComposedLayout, gmem_tiled_copy_QKV: cute.TiledCopy, gmem_tiled_copy_O: cute.TiledCopy, tiled_mma: cute.TiledMma, SharedStorage: cutlass.Constexpr, ): """Kernel function for HSTU attention. :param mQ: query tensor :type mQ: cute.Tensor :param mK: key tensor :type mK: cute.Tensor :param mV: value tensor :type mV: cute.Tensor :param mO: output tensor :type mO: cute.Tensor :param mRAB: RAB(Relative Attention Bias) tensor :type mRAB: cute.Tensor :param sQ_layout: shared memory layout for Q :type sQ_layout: cute.ComposedLayout :param sKV_layout: shared memory layout for K/V :type sKV_layout: cute.ComposedLayout :param sRAB_layout: shared memory layout for RAB :type sRAB_layout: cute.ComposedLayout :param sO_layout: shared memory layout for O :type sO_layout: cute.ComposedLayout :param gmem_tiled_copy_QKV: tiled copy for QKV load :type gmem_tiled_copy_QKV: cute.TiledCopy :param gmem_tiled_copy_O: tiled copy for O store :type gmem_tiled_copy_O: cute.TiledCopy :param tiled_mma: tiled mma :type tiled_mma: cute.TiledMma :param SharedStorage: shared storage :type SharedStorage: cutlass.Constexpr """ # Thread index, block index tidx, _, _ = cute.arch.thread_idx() if cutlass.const_expr(self._enable_block_rasterization): batch_size, num_head, m_block = cute.arch.block_idx() else: m_block, batch_size, num_head = cute.arch.block_idx() # reverse the m_block index m_block = cute.ceil_div(mQ.shape[1], self._m_block_size) - m_block - 1 if cutlass.const_expr(self._is_causal): n_block = ( cute.ceil_div( min((m_block + 1) * self._m_block_size, mK.shape[1]), self._n_block_size, ) - 1 ) # for causal case, only process the first n_block tiles else: n_block = cute.ceil_div(mK.shape[1], self._n_block_size) - 1 # /////////////////////////////////////////////////////////////////////////////// # Get the appropriate tiles for this thread block. # /////////////////////////////////////////////////////////////////////////////// # (m_block_size, head_dim) gQ = cute.local_tile( mQ[batch_size, None, num_head, None], (self._m_block_size, self._head_dim_padded), (m_block, 0), ) # (n_block_size, head_dim, n_block) gK = cute.local_tile( mK[batch_size, None, num_head, None], (self._n_block_size, self._head_dim_padded), (None, 0), ) # (n_block_size, head_dim, n_block) gV = cute.local_tile( mV[batch_size, None, num_head, None], (self._n_block_size, self._head_dim_padded), (None, 0), ) # (m_block_size, n_block_size) gRAB = cute.local_tile( mRAB[batch_size, num_head, None, None], (self._m_block_size, self._n_block_size), (m_block, None), ) # /////////////////////////////////////////////////////////////////////////////// # Get shared memory buffer # /////////////////////////////////////////////////////////////////////////////// smem = cutlass.utils.SmemAllocator() storage = smem.allocate(SharedStorage) sQ = storage.sQ.get_tensor(sQ_layout) sK = storage.sK.get_tensor(sKV_layout) sV = storage.sV.get_tensor(sKV_layout) sRAB = storage.sRAB.get_tensor(sRAB_layout) # Transpose view of V to tensor with layout (head_dim, n_block_size) for tiled mma sVt = cute.composition( sV, cute.make_layout( (self._head_dim_padded, self._n_block_size), stride=(self._n_block_size, 1), ), ) gmem_thr_copy_QKV = gmem_tiled_copy_QKV.get_slice(tidx) # (CPY_Atom, CPY_M, CPY_K) tQgQ = gmem_thr_copy_QKV.partition_S(gQ) tQsQ = gmem_thr_copy_QKV.partition_D(sQ) # (CPY_Atom, CPY_N, CPY_K, n_block) tKgK = gmem_thr_copy_QKV.partition_S(gK) tKsK = gmem_thr_copy_QKV.partition_D(sK) # (CPY_Atom, CPY_N, CPY_K, n_block) tVgV = gmem_thr_copy_QKV.partition_S(gV) tVsV = gmem_thr_copy_QKV.partition_D(sV) # (CPY_Atom, CPY_M, CPY_N, n_block) tRABgRAB = gmem_tiled_copy_QKV.get_slice(tidx).partition_S(gRAB) tRabsRAB = gmem_tiled_copy_QKV.get_slice(tidx).partition_D(sRAB) # /////////////////////////////////////////////////////////////////////////////// # Tile MMA compute thread partitions and allocate accumulators # /////////////////////////////////////////////////////////////////////////////// thr_mma = tiled_mma.get_slice(tidx) tSrQ = thr_mma.make_fragment_A(thr_mma.partition_A(sQ)) tSrK = thr_mma.make_fragment_B(thr_mma.partition_B(sK)) tOrVt = thr_mma.make_fragment_B(thr_mma.partition_B(sVt)) acc_shape_O = thr_mma.partition_shape_C( (self._m_block_size, self._head_dim_padded) ) acc_O = cute.make_rmem_tensor(acc_shape_O, cutlass.Float32) acc_O.fill(0.0) # /////////////////////////////////////////////////////////////////////////////// # Smem copy atom tiling # /////////////////////////////////////////////////////////////////////////////// smem_copy_atom_Q = cute.make_copy_atom( cute.nvgpu.warp.LdMatrix8x8x16bOp(transpose=False, num_matrices=4), self._dtype, ) smem_copy_atom_K = cute.make_copy_atom( cute.nvgpu.warp.LdMatrix8x8x16bOp(transpose=False, num_matrices=4), self._dtype, ) smem_copy_atom_V = cute.make_copy_atom( cute.nvgpu.warp.LdMatrix8x8x16bOp(transpose=True, num_matrices=4), self._dtype, ) smem_copy_atom_RAB = cute.make_copy_atom( cute.nvgpu.warp.LdMatrix8x8x16bOp(transpose=False, num_matrices=4), self._dtype, ) smem_tiled_copy_Q = cute.make_tiled_copy_A(smem_copy_atom_Q, tiled_mma) smem_tiled_copy_K = cute.make_tiled_copy_B(smem_copy_atom_K, tiled_mma) smem_tiled_copy_V = cute.make_tiled_copy_B(smem_copy_atom_V, tiled_mma) smem_tiled_copy_RAB = cute.make_tiled_copy_C(smem_copy_atom_RAB, tiled_mma) smem_thr_copy_Q = smem_tiled_copy_Q.get_slice(tidx) smem_thr_copy_K = smem_tiled_copy_K.get_slice(tidx) smem_thr_copy_V = smem_tiled_copy_V.get_slice(tidx) smem_thr_copy_RAB = smem_tiled_copy_RAB.get_slice(tidx) tSsQ = smem_thr_copy_Q.partition_S(sQ) tSrQ_copy_view = smem_thr_copy_Q.retile(tSrQ) tSsK = smem_thr_copy_K.partition_S(sK) tSrK_copy_view = smem_thr_copy_K.retile(tSrK) tOsVt = smem_thr_copy_V.partition_S(sVt) tOrVt_copy_view = smem_thr_copy_V.retile(tOrVt) tSsRAB = smem_thr_copy_RAB.partition_S(sRAB) # /////////////////////////////////////////////////////////////////////////////// # Predicate: Mark indices that need to copy when problem_shape isn't a multiple # of tile_shape # /////////////////////////////////////////////////////////////////////////////// # Construct identity layout for Q, KV and RAB mcQ = cute.make_identity_tensor(mQ.layout.shape) mcKV = cute.make_identity_tensor(mK.layout.shape) mcRAB = cute.make_identity_tensor(mRAB.layout.shape) cQ = cute.local_tile( mcQ[batch_size, None, num_head, None], (self._m_block_size, self._head_dim_padded), (m_block, 0), ) cKV = cute.local_tile( mcKV[batch_size, None, num_head, None], (self._n_block_size, self._head_dim_padded), (n_block, 0), ) cRAB = cute.local_tile( mcRAB[batch_size, num_head, None, None], (self._m_block_size, self._n_block_size), (m_block, None), ) # Repeat the partitioning with identity layouts tQcQ = gmem_thr_copy_QKV.partition_S(cQ) tKVcKV = gmem_thr_copy_QKV.partition_S(cKV) tRABcRAB = gmem_thr_copy_QKV.partition_S(cRAB) tQpQ = cute.make_rmem_tensor( cute.make_layout( ( tQsQ.shape[0][1], cute.size(tQsQ, mode=[1]), cute.size(tQsQ, mode=[2]), ), stride=(cute.size(tQsQ, mode=[2]), 0, 1), ), cutlass.Boolean, ) tKVpKV = cute.make_rmem_tensor( cute.make_layout( ( tKsK.shape[0][1], cute.size(tKsK, mode=[1]), cute.size(tKsK, mode=[2]), ), stride=(cute.size(tKsK, mode=[2]), 0, 1), ), cutlass.Boolean, ) # Set predicates for head_dim bounds, seqlen_q/k/v bounds is processed at the first tile. for rest_v in cutlass.range_constexpr(tQpQ.shape[0]): for rest_k in cutlass.range_constexpr(tQpQ.shape[2]): tQpQ[rest_v, 0, rest_k] = cute.elem_less( tQcQ[(0, rest_v), 0, rest_k][3], mQ.layout.shape[3] ) for rest_v in cutlass.range_constexpr(tKVpKV.shape[0]): for rest_k in cutlass.range_constexpr(tKVpKV.shape[2]): tKVpKV[rest_v, 0, rest_k] = cute.elem_less( tKVcKV[(0, rest_v), 0, rest_k][3], mK.layout.shape[3] ) # /////////////////////////////////////////////////////////////////////////////// # Prefetch Prologue # /////////////////////////////////////////////////////////////////////////////// # Start async loads of the last mn-tile, where we take care of the mn residue for m in cutlass.range_constexpr(cute.size(tQsQ.shape[1])): if cute.elem_less(tQcQ[0, m, 0][1], mQ.layout.shape[1]): cute.copy( gmem_tiled_copy_QKV, tQgQ[None, m, None], tQsQ[None, m, None], pred=tQpQ[None, m, None], ) else: # Clear the smem tiles to account for predicated off loads tQsQ[None, m, None].fill(0) for n in cutlass.range_constexpr(cute.size(tKsK.shape[1])): if cute.elem_less(tKVcKV[0, n, 0][1], mK.layout.shape[1]): cute.copy( gmem_tiled_copy_QKV, tKgK[None, n, None, n_block], tKsK[None, n, None], pred=tKVpKV[None, n, None], ) else: # Clear the smem tiles to account for predicated off loads tKsK[None, n, None].fill(0) for m in cutlass.range_constexpr(cute.size(tRABcRAB.shape[1])): for n in cutlass.range_constexpr(cute.size(tRABcRAB.shape[2])): if cute.elem_less( tRABcRAB[0, m, n, n_block][1], mRAB.layout.shape[2] ) and cute.elem_less( tRABcRAB[0, m, n, n_block][2], mRAB.layout.shape[3] ): cute.copy( gmem_tiled_copy_QKV, tRABgRAB[None, m, n, n_block], tRabsRAB[None, m, n], ) else: # Clear the smem tiles to account for predicated off loads tRabsRAB[None, m, n].fill(0) cute.arch.cp_async_commit_group() # /////////////////////////////////////////////////////////////////////////////// # Mainloop # /////////////////////////////////////////////////////////////////////////////// for n_block_idx in range(n_block, -1, -1): # wait for smem tile QK before mma caculation for S cute.arch.cp_async_wait_group(0) self.cta_sync_barrier.arrive_and_wait() if n_block_idx == n_block: for n in cutlass.range_constexpr(cute.size(tVsV.shape[1])): if cute.elem_less(tKVcKV[0, n, 0][1], mV.layout.shape[1]): cute.copy( gmem_tiled_copy_QKV, tVgV[None, n, None, n_block_idx], tVsV[None, n, None], pred=tKVpKV[None, n, None], ) else: tVsV[None, n, None].fill(0) else: cute.copy( gmem_tiled_copy_QKV, tVgV[None, None, None, n_block_idx], tVsV[None, None, None], pred=tKVpKV[None, None, None], ) cute.arch.cp_async_commit_group() acc_shape_S = thr_mma.partition_shape_C( (self._m_block_size, self._n_block_size) ) acc_S = cute.make_rmem_tensor(acc_shape_S, cutlass.Float32) rRAB_shape_S = thr_mma.partition_shape_C( (self._m_block_size, self._n_block_size) ) rRAB = cute.make_rmem_tensor(rRAB_shape_S, self._dtype) tSrRAB_copy_view = smem_thr_copy_RAB.retile(rRAB) cute.copy( smem_tiled_copy_RAB, tSsRAB[None, None, None], tSrRAB_copy_view[None, None, None], ) acc_S.store(rRAB.load().to(cutlass.Float32)) # /////////////////////////////////////////////////////////////////////////////// # S gemm calculation # /////////////////////////////////////////////////////////////////////////////// # ldmatrix first QK k-block for mma cute.copy( smem_tiled_copy_Q, tSsQ[None, None, 0], tSrQ_copy_view[None, None, 0], ) cute.copy( smem_tiled_copy_K, tSsK[None, None, 0], tSrK_copy_view[None, None, 0], ) for k in cutlass.range_constexpr(0, cute.size(tSsQ.shape[2])): # ldmatrix next QK k-block for mma if k < cute.size(tSsQ.shape[2]) - 1: cute.copy( smem_tiled_copy_Q, tSsQ[None, None, k + 1], tSrQ_copy_view[None, None, k + 1], ) cute.copy( smem_tiled_copy_K, tSsK[None, None, k + 1], tSrK_copy_view[None, None, k + 1], ) # mma for S=Q@K cute.gemm( tiled_mma, acc_S, tSrQ[None, None, k], tSrK[None, None, k], acc_S, ) # wait for smem tile V for O cute.arch.cp_async_wait_group(0) self.cta_sync_barrier.arrive_and_wait() if n_block_idx > 0: cute.copy( gmem_tiled_copy_QKV, tKgK[None, None, None, n_block_idx - 1], tKsK[None, None, None], pred=tKVpKV[None, None, None], ) # m residue handling for RAB for m in cutlass.range_constexpr(cute.size(tRABcRAB.shape[1])): if cute.elem_less( tRABcRAB[0, m, 0, n_block_idx - 1][1], mRAB.layout.shape[2] ): cute.copy( gmem_tiled_copy_QKV, tRABgRAB[None, m, None, n_block_idx - 1], tRabsRAB[None, m, None], ) else: tRabsRAB[None, m, None].fill(0) cute.arch.cp_async_commit_group() # /////////////////////////////////////////////////////////////////////////////// # silu activation # /////////////////////////////////////////////////////////////////////////////// if self._enable_fast_sigmoid: t1 = acc_S.load() t2 = t1 * 0.5 acc_S.store(t2) for i in cutlass.range_constexpr(cute.size(acc_S.shape[0])): for j in cutlass.range_constexpr(cute.size(acc_S.shape[1])): for k in cutlass.range_constexpr(cute.size(acc_S.shape[2])): ret = llvm.inline_asm( cutlass.Float32.mlir_type, [acc_S[i, j, k].ir_value()], "tanh.approx.f32 $0, $1;", "=f,f", has_side_effects=False, is_align_stack=False, asm_dialect=llvm.AsmDialect.AD_ATT, ) acc_S[i, j, k] = ret t3 = acc_S.load() t4 = t2 * t3 + t2 acc_S.store(t4) else: LOG2_E = 1.4426950408889634074 t1 = acc_S.load() t2 = t1 * -LOG2_E t3 = cute.math.exp2(t2, fastmath=True) + 1.0 t4 = t1 / t3 acc_S.store(t4) mACC = cute.make_identity_tensor( (mRAB.layout.shape[2], mRAB.layout.shape[3]) ) # (seqlen_q, seqlen_kv) cACC = cute.local_tile( mACC[None, None], (self._m_block_size, self._n_block_size), (m_block, n_block_idx), ) if self._is_causal and (n_block - n_block_idx) < cute.ceil_div( self._m_block_size, self._n_block_size ): tACCcACC = thr_mma.partition_C(cACC) for i in cutlass.range_constexpr(cute.size(tACCcACC.shape[0])): for j in cutlass.range_constexpr(cute.size(tACCcACC.shape[1])): for k in cutlass.range_constexpr(cute.size(tACCcACC.shape[2])): if cute.elem_less( tACCcACC[i, j, k][0], tACCcACC[i, j, k][1] ): acc_S[i, j, k] = 0.0 rP = cute.make_rmem_tensor_like(acc_S, self._dtype) rP.store(acc_S.load().to(self._dtype)) # /////////////////////////////////////////////////////////////////////////////// # O gemm calculation # /////////////////////////////////////////////////////////////////////////////// # Convert layout of acc_S to gemm O accept layout. # Due to the mma instruction shape is 16x8x16, we need to convert from (4, MMA_M, MMA_N) to ((4, 2), MMA_M, MMA_N / 2) # (4, MMA_M, MMA_N) -> (4, MMA_M, (2, MMA_N / 2)) rP_layout_divided = cute.logical_divide(rP.layout, (None, None, 2)) rP_mma_view = cute.make_layout( ( (rP_layout_divided.shape[0], rP_layout_divided.shape[2][0]), rP_layout_divided.shape[1], rP_layout_divided.shape[2][1], ), stride=( (rP_layout_divided.stride[0], rP_layout_divided.stride[2][0]), rP_layout_divided.stride[1], rP_layout_divided.stride[2][1], ), ) tOrP = cute.make_tensor(rP.iterator, rP_mma_view) # ldmatrix first V k-block for mma cute.copy( smem_tiled_copy_V, tOsVt[None, None, 0], tOrVt_copy_view[None, None, 0], ) for k in cutlass.range_constexpr(0, cute.size(tOrP.shape[2])): # ldmatrix next V k-block for mma if k < cute.size(tOrP.shape[2]) - 1: cute.copy( smem_tiled_copy_V, tOsVt[None, None, k + 1], tOrVt_copy_view[None, None, k + 1], ) # mma for O=P@V cute.gemm( tiled_mma, acc_O, tOrP[None, None, k], tOrVt[None, None, k], acc_O, ) # /////////////////////////////////////////////////////////////////////////////// # Epilogue # /////////////////////////////////////////////////////////////////////////////// # store acc_O rO = cute.make_rmem_tensor(acc_O.layout, self._dtype) rO.store(acc_O.load().to(self._dtype)) # reuse sQ's data iterator sO_iter = cute.recast_ptr(sQ.iterator, sO_layout.inner) sO = cute.make_tensor(sO_iter, sO_layout.outer) smem_copy_atom_O = cute.make_copy_atom( cute.nvgpu.CopyUniversalOp(), self._dtype ) smem_tiled_copy_O = cute.make_tiled_copy_C(smem_copy_atom_O, tiled_mma) smem_thr_copy_O = smem_tiled_copy_O.get_slice(tidx) taccOrO = smem_thr_copy_O.retile(rO) taccOsO = smem_thr_copy_O.partition_D(sO) # copy acc O from rmem to smem with sts.32(auto vectorization) cute.copy( smem_copy_atom_O, taccOrO, taccOsO, ) gO = cute.local_tile( mO[batch_size, None, num_head, None], (self._m_block_size, self._head_dim_padded), (m_block, 0), ) gmem_thr_copy_O = gmem_tiled_copy_O.get_slice(tidx) tOsO = gmem_thr_copy_O.partition_S(sO) tOgO = gmem_thr_copy_O.partition_D(gO) tOrO = cute.make_fragment_like(tOgO, self._dtype) # sync before all sts are done. self.cta_sync_barrier.arrive_and_wait() # load acc O from smem to rmem for wider vectorization cute.copy( gmem_tiled_copy_O, tOsO, tOrO, ) # predicate for O mcO = cute.make_identity_tensor(mO.layout.shape) cO = cute.local_tile( mcO[batch_size, None, num_head, None], (self._m_block_size, self._head_dim_padded), (m_block, 0), ) tOcO = gmem_thr_copy_O.partition_D(cO) tOpO = cute.make_rmem_tensor( cute.make_layout( (tOgO.shape[0][1], tOgO.shape[1], tOgO.shape[2]), stride=(tOgO.shape[2], 0, 1), ), cutlass.Boolean, ) for rest_v in cutlass.range_constexpr(tOpO.shape[0]): for rest_n in cutlass.range_constexpr(cute.size(tOpO.shape[2])): tOpO[rest_v, 0, rest_n] = cute.elem_less( tOcO[(0, rest_v), 0, rest_n][3], mO.layout.shape[3] ) # copy acc O from rmem to gmem for rest_m in cutlass.range_constexpr(cute.size(tOpO.shape[1])): if cute.elem_less(tOcO[0, rest_m, 0][1], mO.layout.shape[1]): cute.copy( gmem_tiled_copy_O, tOrO[None, rest_m, None], tOgO[None, rest_m, None], pred=tOpO[None, rest_m, None], ) def run_pytorch_hstu_test( dtype: torch.dtype, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, rab: torch.Tensor, 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 """ q = q.to(dtype) k = k.to(dtype) v = v.to(dtype) rab = rab.to(dtype) s_ = torch.matmul(q, k.transpose(-2, -1)) + rab s_ = torch.nn.functional.silu(s_) if is_causal: mask = torch.ones(1, 1, q.shape[2], k.shape[2], dtype=dtype) mask = torch.tril(mask) s_ = s_ * mask.cuda() o = torch.matmul(s_, v).permute(0, 2, 1, 3).contiguous() return o def run( dtype: Type[cutlass.Numeric], batch_size: int, seqlen_q: int, seqlen_kv: int, num_head: int, head_dim: int, m_block_size: int = 128, n_block_size: int = 128, num_threads: int = 128, enable_fast_sigmoid: bool = False, enable_block_rasterization: bool = False, is_causal: bool = False, perf_test: bool = False, **kwargs, ): """ Run the HSTU attention kernel. :param dtype: data type of the input tensors :type dtype: Type[cutlass.Numeric] :param batch_size: batch size :type batch_size: int :param seqlen_q: sequence length of the query :type seqlen_q: int :param seqlen_kv: sequence length of the key :type seqlen_kv: int :param num_head: number of attention heads :type num_head: int :param head_dim: dimension of the head :type head_dim: int :param m_block_size: block size for the m dimension of computation :type m_block_size: int :param n_block_size: block size for the n dimension of computation :type n_block_size: int :param num_threads: number of threads :type num_threads: int :param enable_fast_sigmoid: whether to use fast sigmoid :type enable_fast_sigmoid: bool :param enable_block_rasterization: whether to use block rasterization :type enable_block_rasterization: bool :param is_causal: whether to use causal masking :type is_causal: bool """ assert dtype == cutlass.Float16 or dtype == cutlass.BFloat16 torch_stream = torch.cuda.current_stream() stream = cuda.CUstream(torch_stream.cuda_stream) print("Running Ampere SM80 HSTUAttentionForward test with:") print("batch_size: ", batch_size) print("seqlen_q: ", seqlen_q) print("seqlen_kv: ", seqlen_kv) print("num_head: ", num_head) print("head_dim: ", head_dim) print("m_block_size: ", m_block_size) print("n_block_size: ", n_block_size) print("num_threads: ", num_threads) print("is_causal: ", is_causal) print("enable_fast_sigmoid: ", enable_fast_sigmoid) print("enable_block_rasterization: ", enable_block_rasterization) print("dtype: ", dtype) # reduced tensor num and iter num for functionality test TENSOR_NUM = 1 ITER_NUM = 1 WARMUP_NUM = 0 if perf_test: TENSOR_NUM = 3 ITER_NUM = 100 WARMUP_NUM = 10 # Create tensor Q/K/V/O qs = [ torch.randn( batch_size, seqlen_q, num_head, head_dim, dtype=cutlass_torch.dtype(dtype) ).cuda() for _ in range(TENSOR_NUM) ] ks = [ torch.randn( batch_size, seqlen_kv, num_head, head_dim, dtype=cutlass_torch.dtype(dtype) ).cuda() for _ in range(TENSOR_NUM) ] vs = [ torch.randn( batch_size, seqlen_kv, num_head, head_dim, dtype=cutlass_torch.dtype(dtype) ).cuda() for _ in range(TENSOR_NUM) ] os = [ torch.randn( batch_size, seqlen_q, num_head, head_dim, dtype=cutlass_torch.dtype(dtype) ).cuda() for _ in range(TENSOR_NUM) ] rabs = [ torch.randn( batch_size, num_head, seqlen_q, seqlen_kv, dtype=cutlass_torch.dtype(dtype) ).cuda() for _ in range(TENSOR_NUM) ] fa2_fwd = HSTUAttentionForwardAmpere( dtype, batch_size, seqlen_q, seqlen_kv, num_head, head_dim, m_block_size, n_block_size, num_threads, enable_fast_sigmoid=enable_fast_sigmoid, enable_block_rasterization=enable_block_rasterization, is_causal=is_causal, ) # assume input is 16B align. mqs = [ ( from_dlpack(qs[i], assumed_align=16) .mark_layout_dynamic(leading_dim=3) .mark_compact_shape_dynamic( mode=3, stride_order=qs[i].dim_order(), divisibility=(128 // dtype.width), ) ) for i in range(TENSOR_NUM) ] mks = [ ( from_dlpack(ks[i], assumed_align=16) .mark_layout_dynamic(leading_dim=3) .mark_compact_shape_dynamic( mode=3, stride_order=ks[i].dim_order(), divisibility=(128 // dtype.width), ) ) for i in range(TENSOR_NUM) ] mvs = [ ( from_dlpack(vs[i], assumed_align=16) .mark_layout_dynamic(leading_dim=3) .mark_compact_shape_dynamic( mode=3, stride_order=vs[i].dim_order(), divisibility=(128 // dtype.width), ) ) for i in range(TENSOR_NUM) ] mos = [ ( from_dlpack(os[i], assumed_align=16) .mark_layout_dynamic(leading_dim=3) .mark_compact_shape_dynamic( mode=3, stride_order=os[i].dim_order(), divisibility=(128 // dtype.width), ) ) for i in range(TENSOR_NUM) ] mrabs = [ ( from_dlpack(rabs[i], assumed_align=16) .mark_layout_dynamic(leading_dim=3) .mark_compact_shape_dynamic( mode=3, stride_order=rabs[i].dim_order(), divisibility=(128 // dtype.width), ) ) for i in range(TENSOR_NUM) ] start_event = torch.cuda.Event(enable_timing=True) end_event = torch.cuda.Event(enable_timing=True) kernel = cute.compile( fa2_fwd, mqs[0], mks[0], mvs[0], mos[0], mrabs[0], stream, ) for i in range(0, ITER_NUM): if i == WARMUP_NUM: start_event.record(torch_stream) # Run the kernel kernel( mqs[i % TENSOR_NUM], mks[i % TENSOR_NUM], mvs[i % TENSOR_NUM], mos[i % TENSOR_NUM], mrabs[i % TENSOR_NUM], stream, ) end_event.record(torch_stream) torch.cuda.synchronize(torch_stream) elapsed_time = start_event.elapsed_time(end_event) elapsed_time_avg = elapsed_time / (ITER_NUM - WARMUP_NUM) LAST_USED_TENSOR = (ITER_NUM - 1) % TENSOR_NUM q = qs[LAST_USED_TENSOR].permute(0, 2, 1, 3).contiguous() k = ks[LAST_USED_TENSOR].permute(0, 2, 1, 3).contiguous() v = vs[LAST_USED_TENSOR].permute(0, 2, 1, 3).contiguous() rab = rabs[LAST_USED_TENSOR] kernel_out = os[LAST_USED_TENSOR].cpu() with torch.cuda.stream(torch_stream): ref_bf16 = run_pytorch_hstu_test(torch.bfloat16, q, k, v, rab, is_causal).cpu() ref_fp32 = run_pytorch_hstu_test(torch.float32, q, k, v, rab, is_causal).cpu() torch.cuda.synchronize(torch_stream) assert (kernel_out - ref_fp32).abs().max().item() <= 4 * ( ref_bf16 - ref_fp32 ).abs().max().item() print("Results verified successfully!") if perf_test: print(f"Elapsed time: {elapsed_time_avg:.3f} ms") return elapsed_time_avg * 1000 # return in microseconds if __name__ == "__main__": parser = argparse.ArgumentParser(description="example of HSTU attention with CuTe") parser.add_argument("--dtype", type=cutlass.dtype, default=cutlass.BFloat16) parser.add_argument("--batch_size", type=int, default=4) parser.add_argument("--seqlen_q", type=int, default=2048) parser.add_argument("--seqlen_kv", type=int, default=2048) parser.add_argument("--num_head", type=int, default=4) parser.add_argument("--head_dim", type=int, default=128) parser.add_argument("--m_block_size", type=int, default=64) parser.add_argument("--n_block_size", type=int, default=64) parser.add_argument("--num_threads", type=int, default=128) parser.add_argument( "--no_fast_sigmoid", action="store_false", dest="enable_fast_sigmoid" ) parser.add_argument( "--no_block_rasterization", action="store_false", dest="enable_block_rasterization", ) parser.add_argument("--is_causal", action="store_true", dest="is_causal") parser.add_argument("--perf_test", action="store_true", dest="perf_test") args = parser.parse_args() run( args.dtype, args.batch_size, args.seqlen_q, args.seqlen_kv, args.num_head, args.head_dim, args.m_block_size, args.n_block_size, args.num_threads, args.enable_fast_sigmoid, args.enable_block_rasterization, args.is_causal, args.perf_test, ) print("PASS")