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
dispatch_materialized_mask
def dispatch_materialized_mask[dtype: DType, layout: Layout, //, callback_fn: def[mask_t: MHAMask](mask: mask_t) raises capturing thin -> None](mask_nd: LayoutTensor[dtype, layout], start_pos_nd: OptionalReg[LayoutTensor[DType.uint32, Layout.row_major(Int(-1)), ImmutAnyOrigin]] = None)
Wrap a dense mask tensor in a MaterializedMask and invoke a callback.
Constructs a MaterializedMask from the provided tensor and optional
per-sequence start-position tensor, then calls callback_fn with the
resulting mask. Use this when the mask is provided as an explicit tensor
(e.g. an ALiBi or relative-positional-encoding bias) rather than a
compute-on-the-fly strategy.
Parameters:
- dtype (
DType): Element type of the mask tensor. - layout (
Layout): Layout of the mask tensor. - callback_fn (
def[mask_t: MHAMask](mask: mask_t) raises capturing thin -> None): Parametric callback invoked with theMaterializedMask.
Args:
- mask_nd (
LayoutTensor[dtype, layout]): The mask values tensor with shape(batch, heads, q, k)or compatible broadcast shape. - start_pos_nd (
OptionalReg[LayoutTensor[DType.uint32, Layout.row_major(Int(-1)), ImmutAnyOrigin]]): Optional per-sequence start positions used to offset the key dimension.
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!