# 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 Type import argparse import cuda.bindings.driver as cuda import cutlass 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], 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, 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 :param q: query tensor :param k: key tensor :param v: value tensor :param rab: RAB 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) 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 import torch import cutlass.torch as cutlass_torch 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")