Skip to main content

/

DocsCode

Mojo struct

RaggedMHAOperand

struct RaggedMHAOperand[origin: ImmutOrigin, cache_origin: ImmutOrigin, //, dtype_: DType, layout: Layout, cache_layout: Layout, scale_dtype_: DType = DType.invalid, scale_layout: Layout = Layout()]

An implementation for ragged LayoutTensor arguments to MHA kernels.

Fields​

  • ​buffer (LayoutTensor[RaggedMHAOperand[dtype_, layout, cache_layout, scale_dtype_, scale_layout].dtype, layout, origin]):
  • ​scale_buffer (LayoutTensor[RaggedMHAOperand[dtype_, layout, cache_layout, scale_dtype_, scale_layout].scale_dtype, scale_layout, ImmutAnyOrigin]):
  • ​cache_row_offsets (LayoutTensor[DType.uint32, cache_layout, cache_origin]):

Implemented traits​

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

comptime members​

device_type​

comptime device_type = RaggedMHAOperand[dtype_, layout, cache_layout, scale_dtype_, scale_layout]

dtype​

comptime dtype = dtype_

page_size​

comptime page_size = 0

quantization_enabled​

comptime quantization_enabled = False

quantization_granularity​

comptime quantization_granularity = 0

scale_dtype​

comptime scale_dtype = scale_dtype_

Methods​

__init__​

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

__init__(buffer: LayoutTensor[RaggedMHAOperand[dtype_, layout, cache_layout, scale_dtype_, scale_layout].dtype, layout, origin], scale_buffer: LayoutTensor[RaggedMHAOperand[dtype_, layout, cache_layout, scale_dtype_, scale_layout].scale_dtype, scale_layout, ImmutAnyOrigin], cache_row_offsets: LayoutTensor[DType.uint32, cache_layout, cache_origin]) -> Self

get_type_name​

static get_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) -> UnsafePointer[Scalar[RaggedMHAOperand[dtype_, layout, cache_layout, scale_dtype_, scale_layout].dtype], ImmutAnyOrigin]

Returns:

UnsafePointer

scales_block_paged_ptr​

scales_block_paged_ptr(self, batch_idx: Int, start_tok_idx: Int, head_idx: Int, head_dim_idx: Int = 0) -> UnsafePointer[Scalar[RaggedMHAOperand[dtype_, layout, cache_layout, scale_dtype_, scale_layout].scale_dtype], ImmutAnyOrigin]

Returns:

UnsafePointer

load_scale​

load_scale[width: Int](self, batch_idx: Int, start_tok_idx: Int, head_idx: Int, head_dim_idx: Int) -> SIMD[RaggedMHAOperand[dtype_, layout, cache_layout, scale_dtype_, scale_layout].scale_dtype, width]

Returns:

SIMD

cache_length​

cache_length(self, batch_idx: Int) -> Int

Returns:

Int

max_context_length​

max_context_length(self) -> UInt32

Returns:

UInt32

num_kv_rows​

num_kv_rows(self) -> Int

Returns the total number of tokens in the ragged buffer.

Returns:

Int

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[swizzle_mode: TensorMapSwizzle, *, BN: Int, depth: Int, BK: Int = padded_depth[RaggedMHAOperand[dtype_, layout, cache_layout, scale_dtype_, scale_layout].dtype, swizzle_mode, depth]()](self, ctx: DeviceContext, out tma: TMATensorTile[RaggedMHAOperand[dtype_, layout, cache_layout, scale_dtype_, scale_layout].dtype, 3, _padded_shape[3, RaggedMHAOperand[dtype_, layout, cache_layout, scale_dtype_, scale_layout].dtype, IndexList(VariadicList(BN, 1, BK), Tuple()), swizzle_mode](), _ragged_shape[3, RaggedMHAOperand[dtype_, layout, cache_layout, scale_dtype_, scale_layout].dtype, IndexList(VariadicList(BN, 1, BK), Tuple()), swizzle_mode]()])

Creates a TMA tile for efficient GPU memory transfers.

Returns:

TMATensorTile

create_scale_tma_tile​

create_scale_tma_tile[BMN: Int](self, ctx: DeviceContext, out tma: TMATensorTile[RaggedMHAOperand[dtype_, layout, cache_layout, scale_dtype_, scale_layout].scale_dtype, 2, Index[Int, Int](VariadicPack(1, BMN))])

Returns:

TMATensorTile

create_ragged_tma_tile​

create_ragged_tma_tile[swizzle_mode: TensorMapSwizzle, *, BN: Int, depth: Int, BK: Int = padded_depth[RaggedMHAOperand[dtype_, layout, cache_layout, scale_dtype_, scale_layout].dtype, swizzle_mode, depth]()](self, ctx: DeviceContext, out tma: RaggedTMA3DTile[RaggedMHAOperand[dtype_, layout, cache_layout, scale_dtype_, scale_layout].dtype, swizzle_mode, BN, BK])

Returns:

RaggedTMA3DTile

create_rope_tma_tile​

create_rope_tma_tile[swizzle_mode: TensorMapSwizzle, *, BN: Int, BK: Int, padded_depth: Int](self, ctx: DeviceContext, out tma: TMATensorTile[DType.bfloat16, 3, _padded_shape[3, DType.bfloat16, IndexList(VariadicList(BN, 1, BK), Tuple()), swizzle_mode](), _ragged_shape[3, DType.bfloat16, IndexList(VariadicList(BN, 1, BK), Tuple()), swizzle_mode]()])

Not supported for RaggedMHAOperand.

Returns:

TMATensorTile

create_gather4_tma_tile​

create_gather4_tma_tile[row_width: Int, swizzle_mode: TensorMapSwizzle = TensorMapSwizzle.SWIZZLE_NONE](self, ctx: DeviceContext, out tma: TMATensorTile[RaggedMHAOperand[dtype_, layout, cache_layout, scale_dtype_, scale_layout].dtype, 2, IndexList(VariadicList(4, row_width), Tuple()), IndexList(VariadicList(1, row_width), Tuple())])

Creates a 2D TMA gather4 descriptor for this ragged operand.

Returns:

TMATensorTile

scales_raw_ptr​

scales_raw_ptr(self) -> UnsafePointer[Float32, MutAnyOrigin]

Returns a null pointer. Ragged operands do not support quantization.

Returns:

UnsafePointer

View source

View source

Was this page helpful?

Models

Cookbook

Puzzles

Blog

Community

Contact