IMPORTANT: To view this page as Markdown, append `.md` to the URL (e.g. /max/get-started.md). For the complete documentation index, see llms.txt.
Skip to main content
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

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​

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​

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​

Structs​

  • ​MHAPosition: Position of the MHA-kernel. When decoding=False, q_head_stride == q_num_heads. When decoding=True, q_head_stride == 1.
  • ​NonNullPointer: A pointer with a compile-time guarantee of being non-null.
  • ​NullPointer: A pointer known at compile time to be null, used when an optional pointer argument is absent.
  • ​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.

Traits​

  • ​OptionalPointer: Abstracts over nullable pointers, providing a uniform interface for NonNullPointer and NullPointer.

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.

Edit this page

Was this page helpful?

Models

Cookbook

Puzzles

Blog

Community

Contact