Mojo function
q_tma
q_tma[dtype: DType, //, swizzle_mode: TensorMapSwizzle, *, BM: Int, depth: Int, q_num_heads: Int, group: Int, decoding: Bool, num_qk_stages: Int = 1](ctx: DeviceContext, ptr: UnsafePointer[Scalar[dtype], origin], rows: Int) -> TMATensorTile[dtype, _split_last_layout[dtype](q_smem_shape[dtype, swizzle_mode, BM=BM, group=group, depth=depth, decoding=decoding, num_qk_stages=num_qk_stages](), swizzle_mode, True), _ragged_desc_layout[dtype](q_smem_shape[dtype, swizzle_mode, BM=BM, group=group, depth=depth, decoding=decoding, num_qk_stages=num_qk_stages](), swizzle_mode)]
Returns:
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!