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

scale_and_mask_helper

def scale_and_mask_helper[p_type: DType, p_layout: Layout, mask_t: MHAMask, group: Int, num_n_mmas: Int, WN: Int, MMA_N: Int, simd_width: Int](p_reg_tile: LayoutTensor[p_type, p_layout, address_space=AddressSpace.LOCAL], scale_log2e: Float32, num_keys: Int, bound: Int, lane: Int, warp: Int, mask: mask_t, kv_tile_start_row: Int)

Apply softmax scaling and masking to one P = Q·Kᵀ MMA result tile in registers.

Multiplies each element by scale_log2e and then applies mask to out-of-bounds and masked positions. Only threads with lane < 4 * group carry meaningful data; other threads return immediately. Designed for the decode inner loop where P is a 1-D column-vector across the key dimension.

Parameters:

p_type (DType): Element data type of the P register tile.
p_layout (Layout): Layout of the P register tile.
mask_t (MHAMask): Attention mask type implementing MHAMask.
group (Int): GQA group size (query heads per KV head).
num_n_mmas (Int): Number of MMA operations along the N (key) dimension.
WN (Int): Warp tile width along the N dimension.
MMA_N (Int): MMA instruction width along N.
simd_width (Int): SIMD vector width used in the register tile.

Args:

p_reg_tile (LayoutTensor[p_type, p_layout, address_space=AddressSpace.LOCAL]): Mutable register tile holding Q·Kᵀ values.
scale_log2e (Float32): Pre-multiplied softmax scale (scale * log2e).
num_keys (Int): Total number of valid keys in the sequence.
bound (Int): Inclusive upper bound on the key index for the current tile.
lane (Int): Intra-warp lane ID.
warp (Int): Warp index within the CTA.
mask (mask_t): Mask instance.
kv_tile_start_row (Int): Global key index of the first column in this tile.