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

apply_mask_kernel

def apply_mask_kernel[mask_t: MHAMask, ScoresLayoutType: TensorLayout, scores_origin: MutOrigin, VLLayoutType: TensorLayout, vl_origin: ImmOrigin, CLLayoutType: TensorLayout](output: TileTensor[DType.float32, ScoresLayoutType, scores_origin], valid_length: TileTensor[DType.uint32, VLLayoutType, vl_origin], cache_lengths: TileTensor[DType.uint32, CLLayoutType, ImmutAnyOrigin], mask: mask_t, max_num_keys: Int)

Apply causal mask to the output scores.

Parameters:

mask_t (MHAMask): The MHAMask type applied to each score coordinate.
ScoresLayoutType (TensorLayout): Layout of the output scores tensor.
scores_origin (MutOrigin): Origin of the output scores tensor.
VLLayoutType (TensorLayout): Layout of the valid_length tensor.
vl_origin (ImmOrigin): Origin of the valid_length tensor.
CLLayoutType (TensorLayout): Layout of the cache_lengths tensor.

Args:

output (TileTensor[DType.float32, ScoresLayoutType, scores_origin]): Score matrix with row stride max_num_keys, indexed as [global_seq_idx, key_idx].
valid_length (TileTensor[DType.uint32, VLLayoutType, vl_origin]): Row offsets into output per batch, length batch_size + 1.
cache_lengths (TileTensor[DType.uint32, CLLayoutType, ImmutAnyOrigin]): Per-batch cached-prefix length used to map a local query index to an absolute position.
mask (mask_t): The mask instance applied to each score coordinate.
max_num_keys (Int): Row stride of output and maximum keys per token.