Skip to main content

Mojo function

q_tma

q_tma[dtype: DType, //, swizzle_mode: TensorMapSwizzle, *, BM: Int, depth: Int, q_num_heads: Int, group: Int, decoding: Bool](ctx: DeviceContext, ptr: LegacyUnsafePointer[Scalar[dtype]], rows: Int) -> TMATensorTile[dtype, _split_last_layout[dtype](q_smem_shape[dtype, swizzle_mode, BM=BM, group=group, depth=depth, decoding=decoding](), swizzle_mode, True), _ragged_desc_layout[dtype](q_smem_shape[dtype, swizzle_mode, BM=BM, group=group, depth=depth, decoding=decoding](), swizzle_mode)]

Returns:

TMATensorTile

Was this page helpful?