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

Mojo module

common

Shared NVIDIA GPU attention primitives used by both SM90 and SM100 kernels.

This module hosts helpers that are not architecture-specific so that neither sm90/ nor sm100/ has to import from the other. It currently provides:

• elect(): single-lane election via the elect.sync PTX instruction.

comptime values

ImmutTileTensor1D

comptime ImmutTileTensor1D[dtype: DType] = TileTensor[dtype, Layout[*(), *()], ImmutAnyOrigin]

Parameters

• ​dtype (DType):

KVTMATile

comptime KVTMATile[dtype: DType, swizzle_mode: TensorMapSwizzle, *, BN: Int, BK: Int] = TMATensorTile[dtype, Int(3), _padded_shape[Int(3), dtype, IndexList(BN, Int(1), BK, __list_literal__=NoneType(None)), swizzle_mode](), _ragged_shape[Int(3), dtype, IndexList(BN, Int(1), BK, __list_literal__=NoneType(None)), swizzle_mode]()]

Parameters

• ​dtype (DType):
• ​swizzle_mode (TensorMapSwizzle):
• ​BN (Int):
• ​BK (Int):

QTMATile

comptime QTMATile[dtype: DType, swizzle_mode: TensorMapSwizzle, *, BM: Int, depth: Int, group: Int, decoding: Bool, fuse_gqa: Bool = False, num_qk_stages: Int = Int(1)] = TMATensorTile[dtype, Int(4) if decoding or fuse_gqa else Int(3), _padded_shape[Int(4) if decoding or fuse_gqa else Int(3), dtype, q_smem_shape[dtype, swizzle_mode, BM=BM, group=group, depth=depth, decoding=decoding, fuse_gqa=fuse_gqa, num_qk_stages=num_qk_stages](), swizzle_mode](), _ragged_shape[Int(4) if decoding or fuse_gqa else Int(3), dtype, q_smem_shape[dtype, swizzle_mode, BM=BM, group=group, depth=depth, decoding=decoding, fuse_gqa=fuse_gqa, num_qk_stages=num_qk_stages](), swizzle_mode]()]

Parameters

• ​dtype (DType):
• ​swizzle_mode (TensorMapSwizzle):
• ​BM (Int):
• ​depth (Int):
• ​group (Int):
• ​decoding (Bool):
• ​fuse_gqa (Bool):
• ​num_qk_stages (Int):

Structs

• ​MHAPosition: Position of the MHA-kernel. When decoding=False, q_head_stride == q_num_heads. When decoding=True, q_head_stride == 1.
• ​Pack: Bundles MHA kernel parameters into a single device-passable struct.
• ​PositionSummary: Holds the computed number of keys and the score row index for an attention tile.

Functions

• ​elect: Elects a single lane in the warp via the elect.sync PTX instruction.
• ​get_seq_info: Computes the SeqInfo for the current CTA by querying the transient tile scheduler.
• ​kv_coord: Returns the 3D TMA coordinates for a KV tensor load.
• ​output_reg_to_smem_st_matrix: Stores output register fragments to shared memory using the stmatrix PTX instruction.
• ​q_coord: Returns the coordinates for a tma load on the Q matrix. This load can be 3D, 4D, or 5D.
• ​q_gmem_shape: Computes the global-memory shape for a Q tensor TMA tile based on the tile configuration.
• ​q_smem_shape: Computes the shared-memory shape for a Q tensor TMA tile based on the tile configuration.
• ​q_tma: Creates a split TMA descriptor for the Q tensor, pairing the shared-memory tile shape with the global-memory layout.