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

a_lds_swizzle

def a_lds_swizzle[BK_BYTES: Int]() -> Swizzle

Builds the XOR-16 LDS swizzle for a row-major [BM, BK_BYTES] uint8 A tile.

Returns a Swizzle on the flat in-tile byte offset that XORs the log2(BK_BYTES//16) row bits sitting at flat-bit log2(BK_BYTES) down into the column's 16B-granule bits (base=4), removing A LDS bank conflicts. Applied identically at write (copy_a_tile_to_smem) and read (load_a_frag_from_smem); a mismatch produces wrong logits. BK_BYTES must be a power of two (64/128/256).

Parameters:

BK_BYTES (Int): Width of the A tile in bytes, a power of two.

Returns:

Swizzle: The XOR-16 Swizzle for the A LDS tile.