Skip to main content
Log in

Mojo trait

MHAOperand

This serves as the trait to support arguments to our MHA kernel.

Implemented traits

AnyType, UnknownDestructibility

Aliases

type

alias type

Methods

block_paged_ptr

block_paged_ptr[tile_size: Int](self: _Self, batch_idx: SIMD[uint32, 1], start_tok_idx: SIMD[uint32, 1], head_idx: SIMD[uint32, 1], head_dim_idx: SIMD[uint32, 1] = __init__[__mlir_type.!pop.int_literal](0)) -> UnsafePointer[SIMD[get_vtable_entry(:trait<_nn::_mha_operand::_MHAOperand> _Self, "type"), 1]]

cache_length

cache_length(self: _Self, batch_idx: Int) -> Int

Returns the length of the cache for a given batch index.

max_context_length

max_context_length(self: _Self) -> SIMD[uint32, 1]

Returns the maximum cache length in a given batch index.