For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /get-started.md).

Mojo struct

RelativeLogitsMask

struct RelativeLogitsMask[V: MHAMask, //, visibility: V, dtype_: DType, layout_: Layout, origin_: ImmOrigin]

Causal (optionally sliding-window) mask plus an additive relative-position bias.

The bias is gathered by rel_dist = q_pos - k_pos from a (tokens, heads, extent) table and added on every visible position; distances outside [0, extent) carry no bias. mask() adds the bias before delegating the visibility select to visibility.mask(), so masked lanes come out as exactly MASK_VALUE (never MASK_VALUE + bias).

Parameters

• ​visibility (V): CausalMask() (global) or SlidingWindowCausalMask[window_size]() (local).
• ​dtype_ (DType): Element type of the bias tensor.
• ​layout_ (Layout): Layout of the bias tensor, rank 3 (total_q_tokens, heads, extent).
• ​origin_ (ImmOrigin): Origin of the bias tensor.

Fields

• ​bias (LayoutTensor[dtype_, layout_, origin_]): (total_q_tokens, heads, extent), row r matching q's ragged-flat row.
• ​cache_lengths (LayoutTensor[.uint32, Layout.row_major(Int(-1)), ImmutAnyOrigin]): Cached tokens before this call's new tokens.
• ​input_row_offsets (LayoutTensor[.uint32, Layout.row_major(Int(-1)), ImmutAnyOrigin]): Ragged row offset into bias/q for this call's new tokens, (batch + 1,).

Implemented traits

AnyType, Copyable, Deinitable, DevicePassable, ImplicitlyCopyable, MHAMask, Movable, RegisterPassable, TrivialRegisterPassable

comptime members

apply_log2e_after_mask

comptime apply_log2e_after_mask = True

check_mask_during_decoding

comptime check_mask_during_decoding = True

device_type

comptime device_type = RelativeLogitsMask[visibility, dtype_, layout_, origin_]

mask_out_of_bound

comptime mask_out_of_bound = is_nvidia_gpu()

mask_safe_out_of_bounds

comptime mask_safe_out_of_bounds = True

window_size

comptime window_size = V.sliding_window_size()()

0 means unbounded (plain causal), per this file's convention.

Methods

__init__

def __init__(bias: LayoutTensor[dtype_, layout_, origin_], cache_lengths: LayoutTensor[.uint32, Layout.row_major(Int(-1)), ImmutAnyOrigin], input_row_offsets: LayoutTensor[.uint32, Layout.row_major(Int(-1)), ImmutAnyOrigin]) -> Self

get_type_name

static def get_type_name() -> String

Returns:

String

name

static def name() -> String

Returns:

String

mask

def mask[dtype: DType, width: SIMDLength, //, *, element_type: DType = .uint32](self, coord: IndexList[Int(4), element_type=element_type], score_vec: SIMD[dtype, width]) -> SIMD[dtype, width]

Returns:

SIMD[dtype, width]

status

def status[*, element_type: DType = .uint32](self, seq_id: UInt32, tile_offset: IndexList[Int(2), element_type=element_type], tile_size: IndexList[Int(2), element_type=element_type]) -> TileMaskStatus

Returns:

TileMaskStatus

start_column

def start_column[BM: Int, BN: Int, page_size: Int](self, seq_id: UInt32, row: UInt32) -> UInt32

Returns:

UInt32

start_column_alignment

static def start_column_alignment[BM: Int, BN: Int, page_size: Int]() -> Int

Returns:

Int

total_iters

def total_iters[BM: Int, BN: Int, page_size: Int](self, seq_id: UInt32, row: UInt32, num_cols: UInt32) -> UInt32

Returns:

UInt32

count_nonfull_sets

static def count_nonfull_sets(BM: Int, BN: Int) -> Int

Returns:

Int

last_masked_set_end

def last_masked_set_end[BM: Int, BN: Int, page_size: Int](self, seq_id: UInt32, row: UInt32, num_cols: UInt32) -> UInt32

Returns:

UInt32

masked_set_ends

def masked_set_ends[BM: Int, BN: Int, page_size: Int](self, seq_id: UInt32, row: UInt32, num_cols: UInt32) -> StaticTuple[UInt32, RelativeLogitsMask.count_nonfull_sets(BM, BN)]

Returns:

StaticTuple[UInt32, RelativeLogitsMask.count_nonfull_sets(BM, BN)]

nonfull_sets

static def nonfull_sets[BM: Int, BN: Int]() -> StaticTuple[TileMaskStatus, RelativeLogitsMask.count_nonfull_sets(BM, BN)]

Returns:

StaticTuple[TileMaskStatus, RelativeLogitsMask.count_nonfull_sets(BM, BN)]

mask_strategies

static def mask_strategies[BM: Int, BN: Int]() -> StaticTuple[MaskStrategy, RelativeLogitsMask.count_nonfull_sets(BM, BN)]

Returns:

StaticTuple[MaskStrategy, RelativeLogitsMask.count_nonfull_sets(BM, BN)]

mask_bits

def mask_bits(self, seq_id: UInt32, score_row: Int32, col_start: Int32, num_keys: Int32) -> UInt32

Returns:

UInt32

sliding_window_size

static def sliding_window_size() -> Int

Returns:

Int