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

FixedHeightMinHeap

struct FixedHeightMinHeap[k_dtype: DType, v_dtype: DType, levels: Int]

Maintains a fixed-capacity min-heap of key/value pairs used to track top-k logits.

Parameters

• ​k_dtype (DType): Element type of the heap keys, the token id dtype.
• ​v_dtype (DType): Element type of the heap values, the logit dtype.
• ​levels (Int): Number of heap levels; the heap stores 2**levels - 1 entries.

Fields

• ​k_array (Array[Scalar[k_dtype], FixedHeightMinHeap[k_dtype, v_dtype, levels].num_elements]): Inline array of heap keys, storing token ids.
• ​v_array (Array[Scalar[v_dtype], FixedHeightMinHeap[k_dtype, v_dtype, levels].num_elements]): Inline array of heap values, storing logits.

Implemented traits

AnyType, Deinitable, Movable

comptime members

num_elements

comptime num_elements = ((Int(2) ** levels) - Int(1))

Maximum number of entries the heap can hold.

Methods

__init__

def __init__(out self, *, fill_k: Scalar[k_dtype], fill_v: Scalar[v_dtype])

Initializes the heap with fill values for all slots.

Args:

• ​fill_k (Scalar[k_dtype]): Key value to fill every heap slot with.
• ​fill_v (Scalar[v_dtype]): Value to fill every heap slot with.

swap

def swap(mut self, a: Int, b: Int)

Swaps the key/value pairs at two heap positions.

Args:

• ​a (Int): Index of the first element to swap.
• ​b (Int): Index of the second element to swap.

heap_down

def heap_down(mut self)

Restores the min-heap property by sifting the root down.