For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /max/get-started.md).

Mojo function

kv_coord

def kv_coord[*, depth: Int](row: UInt32, head_idx: UInt32) -> StaticTuple[UInt32, Int(3)]

Returns the 3D TMA coordinates for a KV tensor load.

Parameters:

depth (Int): Head dimension of the attention layer, in elements.

Args:

row (UInt32): Row index along the sequence dimension of the KV tensor.
head_idx (UInt32): KV head index for the tensor load.

Returns:

StaticTuple[UInt32, Int(3)]