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

OffsetPosition

struct OffsetPosition[config: MLA_SM100_Decode_Config, KVLUTType: MHAOperand, ragged: Bool, is_cache_length_accurate: Bool, ValidLengthType: OptionalPointer, decoding_warp_split_k: Bool = False, sparse: Bool = False, has_extra_kv: Bool = False, has_variable_topk: Bool = False]

Computes and stores per-CTA row offsets and KV key ranges for the decode kernel.

Parameters

• ​config (MLA_SM100_Decode_Config): Decode config supplying tile sizes and head counts used to compute Q and output row offsets.
• ​KVLUTType (MHAOperand): MHAOperand providing the KV cache tensor and its cache_length accessor.
• ​ragged (Bool): When True, the valid-lengths tensor is interpreted as input row offsets enabling ragged batching.
• ​is_cache_length_accurate (Bool): When False, the kernel adds the local sequence length to the cache length to compute the total key count.
• ​ValidLengthType (OptionalPointer): OptionalPointer type wrapping the per-batch valid-sequence-length tensor.
• ​decoding_warp_split_k (Bool): When True, the CTA processes a split-K partition of the KV cache (defaults to False).
• ​sparse (Bool): When True, the kernel iterates over a sparse subset of tokens selected by d_indices instead of the full KV cache (defaults to False).
• ​has_extra_kv (Bool): When True, sparse attention additionally attends to a separate extra-KV cache (defaults to False).
• ​has_variable_topk (Bool): When True, the sparse top-k length is read per batch from sparse_topk_lengths instead of using the fixed stride (defaults to False).

Fields

• ​seq_len (Int):
• ​max_seq_len (Int):
• ​num_keys (Int):
• ​q_row_offset (Int):
• ​out_row_offset (Int):
• ​split_idx (Int):
• ​batch_idx (Int):
• ​kv_start_row (Int):
• ​num_keys_this_split (Int):
• ​q_token_idx (Int):

Implemented traits

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

Methods

__init__

def __init__(k: KVLUTType, valid_length: Pointer[Scalar[ValidLengthType.dtype], ImmutAnyOrigin], max_seq_len: Int, num_partitions: Int, batch_size: Int, sparse_indices_stride: Int = Int(0), sparse_topk_lengths: OptionalReg[Pointer[Int32, MutAnyOrigin]] = None, sparse_extra_indices_stride: Int = Int(0), sparse_extra_topk_lengths: OptionalReg[Pointer[Int32, MutAnyOrigin]] = None) -> Self

cache_len

def cache_len(self) -> Int

Returns:

Int

start_pos

def start_pos(self, cache_start_pos: UInt32) -> UInt32

Returns:

UInt32

q_row_offset_at

def q_row_offset_at(self, q_local: Int) -> Int

Returns:

Int

out_row_offset_at

def out_row_offset_at(self, q_local: Int) -> Int

Returns:

Int

q_token_idx_at

def q_token_idx_at(self, q_local: Int) -> Int

Returns:

Int