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]

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:

StaticTuple

View source

Was this page helpful?

Thank you! We'll create more content like this.

Thank you for helping us improve!