Skip to main content

Mojo function

output_reg_to_smem

output_reg_to_smem[BM: Int, BN: Int, WM: Int, depth: Int, kv_type: DType, output_type: DType, accum_type: DType, reg_layout: Layout, o_frag_size: Int, num_consumer_threads: Int, simd_size: Int, swizzle: Swizzle, num_m_mmas: Int, num_consumer: Int, mma_thread_layout: Layout](tid: SIMD[uint32, 1], local_warp_group_idx: SIMD[uint32, 1], warp_x: SIMD[uint32, 1], warp_y: SIMD[uint32, 1], q_smem: UnsafePointer[SIMD[kv_type, 1], address_space=AddressSpace(3), alignment=128], output_reg_tile: LayoutTensor[accum_type, reg_layout, MutableAnyOrigin, address_space=AddressSpace(5)]) -> LayoutTensor[output_type, row_major(BM, depth), MutableAnyOrigin, address_space=AddressSpace(3)]

Returns:

LayoutTensor

Was this page helpful?