IMPORTANT: To view this page as Markdown, append `.md` to the URL (e.g. /max/get-started.md). For the complete documentation index, see llms.txt.
Skip to main content
For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /max/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, ImplicitlyCopyable, ImplicitlyDeletable, Movable, RegisterPassable, TrivialRegisterPassable

Methods​

__init__​

def __init__(k: KVLUTType, valid_length: Pointer[Scalar[ValidLengthType.dtype], ImmutAnyOrigin, _safe=False], max_seq_len: Int, num_partitions: Int, batch_size: Int, sparse_indices_stride: Int = Int(0), sparse_topk_lengths: OptionalReg[Pointer[Int32, MutAnyOrigin, _safe=False]] = None, sparse_extra_indices_stride: Int = Int(0), sparse_extra_topk_lengths: OptionalReg[Pointer[Int32, MutAnyOrigin, _safe=False]] = 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

Edit this page

Was this page helpful?

Models

Cookbook

Puzzles

Blog

Community

Contact