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 Types

enum class Mode : uint8_t

Row validation policy.

Values:

enumerator kDense

Live mappings form one dense prefix followed by sentinels.

enumerator kSparseWindow

Live mappings may occupy disjoint logical ranges, as an SWA window advances.

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

KVPageTable(

int32_t maxBatch,

int32_t maxPagesPerSeq,

int32_t numPages,

Mode mode

)

Construct a page table with an explicit row validation mode.

Throws:

std::runtime_error – if any argument is not positive or mode is unsupported

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 setEntry(int32_t slot, int32_t logicalPage, int32_t kPageId)

Set one logical mapping and its derived V mapping.

Throws:

std::runtime_error – if slot, logicalPage, or kPageId is out of range

void clearEntry(int32_t slot, int32_t logicalPage)

Clear one logical K/V mapping to the unused-page sentinel.

Throws:

std::runtime_error – if slot or logicalPage is out of range

bool checkInvariants(std::string &error) const

Validate the host table: every K id is the sentinel or in [0, numPages), and every V id is derived from its K id. Dense rows reject live mappings after a sentinel; sparse-window rows instead reject duplicate live K ids within a slot.

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

void uploadDirty(cudaStream_t stream)

Upload only K/V entries changed since the previous upload. Adjacent dirty entries are coalesced into one H2D copy. Sparse-window mutations maintain their invariants incrementally, so this operation scales with the number of changed entries rather than the logical table width. On the first dirty upload, the device table is initialized to sentinels before applying the changes.

Throws:

std::runtime_error – if a dense table is invalid or a copy fails

inline size_t lastUploadEntryCount() const

Number of int32 K/V entries copied by the most recent successful upload operation.

inline size_t lastUploadRangeCount() const

Number of H2D ranges submitted by the most recent successful upload operation.

void gatherRows(

rt::Tensor &destination,

rt::Tensor const &residentSlots,

int32_t numRows,

cudaStream_t stream

) const

Gather resident rows into an active-step table without changing resident ownership or moving KV payload. destination must have fixed storage for at least [numRows, 2, maxPagesPerSeq] INT32 values.

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 maxBatch() const

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 a row or entry mutator since). Conservative: any mutator call clears this even if the supplied ids happen to describe an identity mapping, so callers that gate identity-only consumers 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 = {}