Skip to main content

/

DocsCode

Mojo struct

RaggedMHAOperand

@register_passable(trivial) struct RaggedMHAOperand[dtype_: DType, layout: Layout, cache_layout: Layout]

An implementation for ragged NDBuffer arguments to MHA kernels.

Fields​

  • ​buffer (LayoutTensor[RaggedMHAOperand[dtype_, layout, cache_layout].dtype, layout, MutAnyOrigin]):
  • ​cache_row_offsets (LayoutTensor[DType.uint32, cache_layout, MutAnyOrigin]):

Implemented traits​

AnyType, Copyable, DevicePassable, ImplicitlyCopyable, ImplicitlyDestructible, MHAOperand, Movable

comptime members​

__copyinit__is_trivial​

comptime __copyinit__is_trivial = True

__del__is_trivial​

comptime __del__is_trivial = True

__moveinit__is_trivial​

comptime __moveinit__is_trivial = True

device_type​

comptime device_type = RaggedMHAOperand[dtype_, layout, cache_layout]

dtype​

comptime dtype = dtype_

page_size​

comptime page_size = 0

Methods​

__init__​

__init__(buffer: LayoutTensor[RaggedMHAOperand[dtype_, layout, cache_layout].dtype, layout, MutAnyOrigin], cache_row_offsets: LayoutTensor[DType.uint32, cache_layout, MutAnyOrigin]) -> Self

get_type_name​

static get_type_name() -> String

Returns:

String

get_device_type_name​

static get_device_type_name() -> String

Returns:

String

block_paged_ptr​

block_paged_ptr[tile_size: Int](self, batch_idx: UInt32, start_tok_idx: UInt32, head_idx: UInt32, head_dim_idx: UInt32 = 0) -> LegacyUnsafePointer[Scalar[RaggedMHAOperand[dtype_, layout, cache_layout].dtype]]

Returns:

LegacyUnsafePointer

cache_length​

cache_length(self, batch_idx: Int) -> Int

Returns:

Int

max_context_length​

max_context_length(self) -> UInt32

Returns:

UInt32

row_idx​

row_idx(self, batch_idx: UInt32, start_tok_idx: UInt32) -> UInt32

Returns the row idx when viewing the memory as a matrix.

Returns:

UInt32

create_tma_tile​

create_tma_tile[BN: Int, depth: Int, swizzle_mode: TensorMapSwizzle, BK: Int = depth](self, ctx: DeviceContext, out tma: TMATensorTile[RaggedMHAOperand[dtype_, layout, cache_layout].dtype, _split_last_layout[RaggedMHAOperand[dtype_, layout, cache_layout].dtype](IndexList[3, DType.int64](BN, 1, BK, Tuple[]()), swizzle_mode, True), _ragged_desc_layout[RaggedMHAOperand[dtype_, layout, cache_layout].dtype](IndexList[3, DType.int64](BN, 1, BK, Tuple[]()), swizzle_mode)])

Creates a TMA tile for efficient GPU memory transfers.

Returns:

TMATensorTile

View source

View source

Was this page helpful?

Models

Cookbook

Puzzles

Blog

Community

Contact