veloxquant_mlx.dsa package

Hand-implemented data structures backing VeloxQuant-MLX’s quantizer internals.

Groups the manually written (no heapq/stdlib shortcuts) structures used by the codebook, outlier-tracking, and pipeline-validation code: AVLTree / VoronoiTree for nearest-centroid search, BitPackBuffer for sub-byte index packing, QuantizationGraph for DAG validation and topological ordering of handler chains, MaxHeap / SortedChannelIndex for streaming top-k channel tracking, and RingBuffer for fixed-capacity FIFO windows.

class veloxquant_mlx.dsa.AVLNode(key, value)[source]

Bases: object

Single node of an AVL tree.

Variables:

  • key (float) – Float key (centroid value).

  • value (int) – Integer payload (centroid index in the codebook).

  • left (AVLNode | None) – Left child.

  • right (AVLNode | None) – Right child.

  • height (int) – Cached subtree height (1 for a leaf).

Parameters:
key: float
value: int
left: AVLNode | None
right: AVLNode | None
height: int
class veloxquant_mlx.dsa.AVLTree[source]

Bases: object

Self-balancing AVL binary search tree.

Keys are float centroid values; values are integer centroid indices. All rotations are implemented manually to satisfy the balance invariant: |balance_factor(node)| <= 1 for every node.

Parameters:

empty. (None. Tree starts)

root: AVLNode | None
insert(key, value)[source]

Insert a (key, value) pair into the tree.

Parameters:

  • key (float) – Float key.

  • value (int) – Integer index payload.

Return type:

None

search_nearest(query)[source]

Find the (key, value) pair with key closest to query.

Parameters:

query (float) – Target value.

Return type:

tuple[float, int]

Returns:

Tuple (nearest_key, nearest_value).

Raises:

RuntimeError – If the tree is empty.

range_query(lo, hi)[source]

Return all values whose keys fall in [lo, hi].

Parameters:

  • lo (float) – Inclusive lower bound.

  • hi (float) – Inclusive upper bound.

Return type:

list[int]

Returns:

List of integer values with keys in [lo, hi].

class veloxquant_mlx.dsa.VoronoiTree[source]

Bases: object

Nearest-centroid search structure wrapping AVLTree.

Falls back to linear scan (numpy.argmin) for small codebooks (n_centroids <= LINEAR_THRESHOLD).

Parameters:

build (None. Call)

LINEAR_THRESHOLD: int = 16
build(centroids)[source]

Populate the search structure from a 1-D centroid array.

Parameters:

centroids (ndarray) – 1-D float array of centroid values, shape (k,).

Return type:

None

nearest(value)[source]

Return the index of the centroid nearest to value.

Parameters:

value (float) – Query float.

Return type:

int

Returns:

Index of the nearest centroid.

Raises:

RuntimeError – If build() has not been called.

class veloxquant_mlx.dsa.BitPackBuffer(b)[source]

Bases: object

Compact buffer that packs b-bit unsigned integers into uint8 bytes.

Supports b ∈ {1, 2, 3, 4}. All bit manipulation is implemented manually.

For non-power-of-2 bit widths (b=3), 8 values fit exactly in 3 bytes: each group of 8 values is packed into 3 bytes (24 bits / 3 bits = 8 values).

Parameters:

b (int) – Bits per element. Must be one of {1, 2, 3, 4}.

Raises:

ValueError – If b is not in {1, 2, 3, 4}.

pack(indices)[source]

Pack an array of b-bit unsigned integers into bytes.

Parameters:

indices (ndarray) – Array of uint8 values in range [0, 2^b). Shape (n,).

Return type:

ndarray

Returns:

Packed uint8 array. Length is ceil(n * b / 8).

Raises:

ValueError – If any value exceeds 2^b - 1.

unpack(packed, n)[source]

Unpack a byte array back into n b-bit unsigned integers.

Parameters:

  • packed (ndarray) – Packed uint8 array produced by pack().

  • n (int) – Original number of values.

Return type:

ndarray

Returns:

Uint8 array of length n with values in [0, 2^b).

class veloxquant_mlx.dsa.QuantizationGraph[source]

Bases: object

Directed Acyclic Graph of QuantizationHandler nodes.

Used to validate user-assembled pipelines (detect cycles) and to determine a correct execution order via topological sort (Kahn’s algorithm).

Nodes are QuantizationHandler instances identified by their id(). Edges represent data-flow dependencies (src must execute before dst).

Parameters:

empty. (None. Graph starts)

add_node(handler)[source]

Register a handler node.

Parameters:

handler (Any) – A QuantizationHandler instance.

Return type:

None

add_edge(src, dst)[source]

Add a directed edge src → dst.

Both nodes must have been added with add_node() first.

Parameters:

  • src (Any) – Source handler (must execute before dst).

  • dst (Any) – Destination handler.

Raises:

KeyError – If either node is not registered.

Return type:

None

has_cycle()[source]

Return True if the graph contains at least one directed cycle.

Uses Kahn’s algorithm: if not all nodes can be processed, a cycle exists.

Return type:

bool

Returns:

True if a cycle is present, False if the graph is a DAG.

topological_sort()[source]

Return nodes in topological order using Kahn’s algorithm.

Kahn’s algorithm:

  1. Compute in-degree for each node.

  2. Enqueue all nodes with in-degree 0.

  3. Repeatedly dequeue a node, append to result, and decrement in-degrees of its successors. Enqueue successors that reach 0.

  4. If |result| < |nodes|, a cycle exists.

Return type:

list[Any]

Returns:

List of handler objects in valid execution order.

Raises:

CyclicPipelineError – If the graph contains a cycle.

critical_path()[source]

Return the longest path through the DAG (by number of nodes).

Uses the topological order computed by Kahn’s algorithm and a dynamic programming pass: dist[v] = max(dist[u] + 1) for all u→v.

Return type:

list[Any]

Returns:

List of handler objects forming the critical (longest) path.

Raises:

CyclicPipelineError – If the graph contains a cycle.

class veloxquant_mlx.dsa.MaxHeap[source]

Bases: Generic[T]

Binary max-heap backed by an internal array.

Nodes are (priority, value) pairs. The heap property is priority[parent] >= priority[child] for all nodes.

All heap operations are implemented manually without the heapq module. value is typically an int (a channel/token index), but SortedChannelIndex uses (channel_idx, version) tuples, hence the generic parameter rather than a fixed int.

Parameters:

empty. (None. The heap starts)

push(priority, value)[source]

Insert a (priority, value) pair.

Parameters:

  • priority (float) – Comparison key (higher = higher priority).

  • value (TypeVar(T)) – Associated payload.

Return type:

None

pop()[source]

Remove and return the maximum (priority, value) pair.

Return type:

tuple[float, TypeVar(T)]

Returns:

Tuple of (priority, value).

Raises:

IndexError – If the heap is empty.

peek()[source]

Return but do not remove the maximum pair.

Raises:

IndexError – If the heap is empty.

Return type:

tuple[float, TypeVar(T)]

class veloxquant_mlx.dsa.SortedChannelIndex[source]

Bases: object

Max-heap–based structure for tracking top-k high-magnitude channels.

Used by OutlierDetector to maintain a running top-k across streaming key vectors during the prefill phase.

Each entry maps a channel index to a magnitude estimate. The structure supports O(log n) insert, O(k log n) top-k query, and O(log n) update.

Parameters:

empty. (None. The index starts)

insert(channel_idx, magnitude)[source]

Insert or update a channel’s magnitude.

If the channel already exists, the magnitude is updated (old entry is lazily invalidated). Re-inserting the same magnitude for a channel that already holds it is a no-op — this keeps the heap bounded at O(number of distinct channels observed) instead of growing by one entry per call.

Parameters:

  • channel_idx (int) – Index of the channel (0-indexed).

  • magnitude (float) – Non-negative magnitude value.

Return type:

None

update(channel_idx, new_magnitude)[source]

Update the magnitude of an existing channel.

Internally uses lazy deletion — the old entry remains in the heap but is skipped when popped because the stored magnitude no longer matches _valid.

Parameters:

  • channel_idx (int) – Channel whose magnitude should be updated.

  • new_magnitude (float) – Replacement magnitude value.

Return type:

None

top_k(k)[source]

Return indices of the k channels with the highest magnitudes.

This operation does not remove elements from the heap.

Parameters:

k (int) – Number of top channels to return.

Return type:

list[int]

Returns:

List of up to k channel indices, ordered from highest to lowest magnitude.

class veloxquant_mlx.dsa.RingBuffer(capacity)[source]

Bases: Generic[T]

Fixed-capacity circular buffer with O(1) append and index access.

When full, the oldest element is silently overwritten (FIFO eviction). Supports negative indexing (buf[-1] is the most recently appended item).

Internal layout:

_data: list of length _capacity, pre-filled with None. _head: index in _data of the oldest element. _size: number of currently live elements.

Parameters:

capacity (int) – Maximum number of elements the buffer can hold. Must be >= 1.

Raises:

ValueError – If capacity < 1.

append(item)[source]

Append an item to the buffer.

If the buffer is full the oldest item is evicted and returned.

Parameters:

item (TypeVar(T)) – Element to append.

Return type:

Optional[TypeVar(T)]

Returns:

The evicted element if the buffer was full, otherwise None.

is_full()[source]

Return True if the buffer is at capacity.

Return type:

bool

to_list()[source]

Return all elements as a plain list (oldest first).

Return type:

list[TypeVar(T)]

Returns:

List of length len(self).

to_stacked()[source]

Stack all stored MLX arrays along axis 0.

Only valid when T = mx.array. Uses mx.stack for efficiency.

Return type:

array

Returns:

Single MLX array of shape (size, *element_shape).

Raises:

RuntimeError – If the buffer is empty.

Submodules