Skip to main content
Log in

Mojo struct

NDBufferMHAOperand

@register_passable(trivial) struct NDBufferMHAOperand[type_: DType, rank: Int, shape: DimList, stride: DimList]

An implementation for NDBuffer arguments to MHA kernels.

Aliases

  • type = type_:

Fields

  • buffer (NDBuffer[type_, rank, MutableAnyOrigin, shape, stride]):

Implemented traits

AnyType, MHAOperand, UnknownDestructibility

Methods

__init__

__init__(buffer: NDBuffer[type_, rank, MutableAnyOrigin, shape, stride]) -> Self

block_paged_ptr

block_paged_ptr[tile_size: Int](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[type_, 1]]

cache_length

cache_length(self, batch_idx: Int) -> Int

max_context_length

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

Was this page helpful?