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 module

warp_spec_matmul

AMD Warp-Specialized Matrix Multiplication.

Architecture Overview:

• Producer warps: Load tiles from global to shared memory
  • A producers: Load M×K tiles from matrix A
  • B producers: Load N×K tiles from matrix B
• Consumer warps: Perform matrix multiplication using shared memory tiles
• Ring buffer: Coordinates producer-consumer synchronization with barriers

Data Flow:

1. Producers load tiles into shared memory stages
2. Barriers ensure data is ready before consumers access it
3. Consumers compute partial results and accumulate
4. Final results written back to global memory

Memory Layout:

• Shared memory is divided into pipeline stages for overlapping
• Each stage contains block tiles that are further divided into warp tiles
• Swizzling may be applied to avoid bank conflicts

Ring Buffer Configuration:

• Uses SingleCounterSync strategy by default (single atomic counter per tile)
• Can be changed to SplitCounterSync in the RingBuffer type aliases for reduced contention
• The trait-based design allows easy experimentation with different sync strategies

comptime values

GlobalTensor

comptime GlobalTensor[dtype: DType, layout: Layout] = LayoutTensor[dtype, layout, MutAnyOrigin, address_space=AddressSpace.GLOBAL]

Parameters

• ​dtype (DType):
• ​layout (Layout):

Functions

• ​determine_thread_role: Returns (role, consumer_warp_id within role group).
• ​get_producer_warp_thread_layout: Computes the thread-to-element layout used by a producer warp when loading a tile.
• ​lgkm_wait: Emits an s_waitcnt lgkmcnt(0) instruction to wait for outstanding shared memory loads.
• ​run_producer: Generic producer function for loading matrix tiles from global to shared memory.
• ​smem_tile_layout: Builds the shared memory layout for a matrix tile used by the warp-specialized kernel.
• ​validate_config: Validates the configuration parameters for the matrix multiplication kernel.
• ​warp_specialized_matmul: Enqueues the warp-specialized matrix multiplication kernel onto the device.
• ​warp_specialized_matmul_kernel: Runs the warp-specialized matrix multiplication kernel on an AMD GPU.

Was this page helpful?

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

Thank you for helping us improve!

😔 What went wrong?

Some code doesn’t work

It includes inaccurate information

It's missing information I need

It was difficult to understand

Other

Submit