KV Page Table

class KVPageTable

One logical KV page table per model (shared by the K and V halves; V page id == K page id + numPages) plus its derived K/V kernel view.

Host state is the source of truth: mHost stores, per slot, the K page ids (-1 for unused); the V id of a live entry is always k + numPages and -1 stays -1. kernelView() is the single device tensor consumed by the paged kernels, laid out as [maxBatch, 2, maxPagesPerSeq] with the K half first and the derived V half second.

Public Functions

KVPageTable(

int32_t maxBatch,

int32_t maxPagesPerSeq,

int32_t numPages

)

Construct a page table for up to maxBatch slots of maxPagesPerSeq logical pages each, backed by a pool of numPages physical pages.

Throws:

std::runtime_error – if any argument is not positive

~KVPageTable() noexcept

KVPageTable(KVPageTable const&) = delete

KVPageTable &operator=(KVPageTable const&) = delete

void setIdentity()

Assign every slot its static identity range: slot b gets K pages [b*maxPagesPerSeq, (b+1)*maxPagesPerSeq) (V = K + numPages).

void setRow(int32_t slot, int32_t const *kPageIds, int32_t count)

Set slot slot’s live K page ids from kPageIds[0..count); V ids are derived as k + numPages. Entries [count, maxPagesPerSeq) are cleared.

Throws:

std::runtime_error – if the row description or a page id is invalid

void setRows(std::vector<KVPageTableRowUpdate> const &updates)

Prevalidate all row updates, then apply them as one host-side commit. A slot may appear at most once. An empty row clears that slot.

Throws:

std::runtime_error – without changing any row if any update is invalid

void compactRows(

std::vector<int32_t> const &oldToNew,

int32_t newBatch

)

Compact logical slots according to one immutable old-to-new mapping. Every destination in [0, newBatch) must appear exactly once; -1 retires an old slot. This changes page-table rows only and never copies KV data.

Throws:

std::runtime_error – without changing any row if the mapping is invalid

bool checkInvariants(std::string &error) const

Validate the host K table: every id is either the sentinel -1 or in [0, numPages), and no live id follows a sentinel within a row.

Parameters:

error – Set to a description of the first violation found

Returns:

true if the table is valid

bool upload(cudaStream_t stream)

Validate the table and enqueue copies for coalesced dirty-row ranges. The first call uploads the complete table. A later call with no dirty rows performs no CUDA operation and does not wait for a prior upload.

Throws:

std::runtime_error – if checkInvariants fails or the copy fails

Returns:

true if at least one device copy was enqueued; false for a no-op

rt::Tensor const &kernelView() const

The device tensor consumed by the paged kernels: int32 [maxBatch, 2, maxPagesPerSeq].

rt::Tensor &kernelView()

Mutable overload for binding into a TensorMap (which stores non-owning Tensor*).

int32_t const *hostRow(int32_t slot) const

Host K row for slot slot (maxPagesPerSeq entries); used by writeKV/gather helpers that need the logical page ids on host.

inline int32_t maxPagesPerSeq() const

Row stride of kernelView() (the maxPagesPerSeq this table was constructed with).

inline int32_t numPages() const

inline bool isIdentity() const

True only if the table’s current contents were last set by setIdentity() (and never touched by setRow() since). Conservative: any setRow() call clears this even if the supplied ids happen to describe an identity mapping, so callers that gate identity-only consumers (e.g. the Alpamayo action runner) fail closed rather than risk a false positive.

struct KVPageTableRowUpdate

One non-owning host-row update consumed synchronously by KVPageTable::setRows(). kPageIds may be null only when count is zero.

Public Members

int32_t slot = {}

int32_t const *kPageIds = {}

int32_t count = {}