Mojo function
q_coord
q_coord[*, depth: Int, swizzle_granularity: Int, decoding: Bool](row: UInt32, head_idx: UInt32) -> StaticTuple[UInt32, (4 if decoding else 3 + Int.__init__[Bool](_should_split_last_dim(depth, swizzle_granularity)))]
Returns the coordinates for a tma load on the Q matrix. This load can be 3D, 4D, or 5D.
Arguments: row: the row to load from. head_idx: q_head_idx if prefill, kv_head_idx if decoding.
Returns:
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!