Mojo function
q_out_tma
q_out_tma[dtype: DType, //, BM: Int, depth: Int, swizzle_mode: TensorMapSwizzle, *, q_num_heads: Int, decoding: Bool](ctx: DeviceContext, ptr: UnsafePointer[SIMD[dtype, 1]], rows: Int, out res: TMATensorTile[dtype, tile_layout_k_major[::DType,::Int,::Int,::TensorMapSwizzle](), _tma_desc_tile_layout[::DType,::Int,::IndexList[$1, ::DType()])
Returns:
TMATensorTile
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!