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

AndMask

struct AndMask[T: MHAMask, S: MHAMask, //, lhs: T, rhs: S]

Mask that's the AND of two masks. If both masks mask off an element, the element is masked off.

Implemented traits

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

comptime members

apply_log2e_after_mask

comptime apply_log2e_after_mask = T.apply_log2e_after_mask or S.apply_log2e_after_mask

check_mask_during_decoding

comptime check_mask_during_decoding = T.check_mask_during_decoding and S.check_mask_during_decoding

device_type

comptime device_type = AndMask[lhs, rhs]

mask_out_of_bound

comptime mask_out_of_bound = T.mask_out_of_bound or S.mask_out_of_bound

mask_safe_out_of_bounds

comptime mask_safe_out_of_bounds = T.mask_safe_out_of_bounds and S.mask_safe_out_of_bounds

Methods

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 = 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 = 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, AndMask.count_nonfull_sets(BM, BN)]

Returns:

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

nonfull_sets

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

Returns:

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

mask_strategies

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

Returns:

StaticTuple[MaskStrategy, AndMask.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