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
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 stores2**levels - 1entries.
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
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:
heap_down
def heap_down(mut self)
Restores the min-heap property by sifting the root down.
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!