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_pack_o_row

def scale_pack_o_row[n: Int, //, output_type: DType, w: Int, start: Int = Int(0)](o_vals: Array[Float32, n], inv_row_sum: Float32) -> SIMD[DType.uint32, (w // Int(2))]

Scale the w f32 O lanes o_vals[start : start + w] by inv_row_sum, cast to the 2-byte output_type, and pack into w // 2 u32 lanes (the row-major 16 B SWIZZLE_NONE store register).

o_vals is a tcgen05_ld result; start/w window it so one wide TMEM load can feed several stores (depth512 loads 16 lanes, stores two 8-lane blocks). Compute stays in f32x2 (64-bit) chunks because LLVM scalarizes wider SIMD here; only the packed u32 store register is built wide. Shared by the SM100 O-store writeback helpers (fa4_scale_write_output, depth512_scale_write_output).

Parameters:

n (Int): Total number of f32 O lanes in o_vals (inferred).
output_type (DType): Target 2-byte dtype (bf16/f16) to cast the scaled lanes to.
w (Int): Number of f32 lanes to scale and pack; must equal the width of one 16 B SWIZZLE_NONE block.
start (Int): Starting index into o_vals for the window (defaults to 0).

Args:

o_vals (Array[Float32, n]): tcgen05_ld result holding the f32 O lanes; the window [start, start + w) is scaled and packed.
inv_row_sum (Float32): Inverse of the softmax row sum, multiplied into each lane to normalize the output.

Returns:

SIMD[DType.uint32, (w // Int(2))]