Class KvCachePool

java.lang.Object
smile.llm.cache.KvCachePool
All Implemented Interfaces:
AutoCloseable
public class KvCachePool extends Object implements AutoCloseable
Physical KV-cache memory pool backed by a RadixCache for prefix-sharing across requests.

On CUDA, sizing follows SGLang --mem-fraction-static (y): 

  staticBudget   = totalGpuMemory × y          // weights + static pools + KV
  dynamicReserve = totalGpuMemory × (1 − y)    // activations / temps (left free)
  kvBudget       = staticBudget − memoryAlreadyUsed
  numSlots       = min(kvBudget / bytesPerToken, maxBatchSize × maxSeqLen)
Call allocate(KvCacheLayout, Device, ScalarType, double, int) only after model weights (and any other static pools such as DeltaNet state) are on device so memoryAlreadyUsed is correct.

Slot indices are managed as fixed-size pages. The embedded RadixCache maps token prefixes to those indices so that shared system prompts and conversation history can be reused without recomputation. Attention layers write and read activations through put(int, int, Tensor, Tensor) / get(int, int) using a per-request slot map established by bindRequests(int, int) or bindWithPrefix(int[], int).

The pool is static: bind claims free pages only and never grows the underlying buffers. When the requested capacity exceeds free pages (after radix eviction), bind clamps to what is available. Generation should stop at requestCapacity() and return partial results rather than allocating more KV.

See Also:

  • Field Details

    • DEFAULT_PAGE_SIZE

      public static final int DEFAULT_PAGE_SIZE
      Default page size (tokens per allocation unit).
      See Also:

  • Constructor Details

    • KvCachePool

      public KvCachePool(int numLayers, int numSlots, int numKvHeads, int headDim, int pageSize, Device device, ScalarType dtype)
      Constructor.
      Parameters:
      numLayers - number of transformer layers.
      numSlots - total token slots (must be a multiple of pageSize).
      numKvHeads - number of KV heads.
      headDim - head embedding dimension.
      pageSize - tokens per page.
      device - compute device.
      dtype - element dtype.

  • Method Details

    • kScale

      public float kScale()
      Returns the K-scale used for FP8 store.
      Returns:
      K-scale used for FP8 store (1.0 when not FP8).

    • vScale

      public float vScale()
      Returns the V-scale used for FP8 store.
      Returns:
      V-scale used for FP8 store (1.0 when not FP8).

    • computeDtype

      public ScalarType computeDtype()
      Returns the dtype used when dequantizing FP8 KV for attention.
      Returns:
      dtype used when dequantizing FP8 KV for attention.

    • allocate

      public static KvCachePool allocate(KvCacheLayout layout, Device device, ScalarType dtype, double memFraction, int pageSize)
      Allocates a pool using SGLang mem-fraction-static semantics.

      Call this after model weights (and other static buffers) are loaded. On CUDA, memFraction is applied to total device memory; KV receives whatever remains inside that static region.

      Parameters:
      layout - family-agnostic cache layout.
      device - compute device.
      dtype - cache element dtype (typically the model weight dtype).
      memFraction - static-region fraction of total GPU memory ((0, 1]).
      pageSize - tokens per page.
      Returns:
      the allocated pool.

    • computeStaticKvBudget

      public static KvCachePool.StaticKvBudget computeStaticKvBudget(long total, long free, double memFraction, long bytesPerToken, int pageSize, int maxBatchSize, int maxSeqLen)
      Computes SGLang-style KV slot budget from device memory readings.
      Parameters:
      total - total device memory bytes.
      free - free device memory bytes (after emptyCache).
      memFraction - static-region fraction y.
      bytesPerToken - bytes for one token of K+V across all layers.
      pageSize - tokens per page.
      maxBatchSize - configured max batch.
      maxSeqLen - configured max sequence length.
      Returns:
      budget breakdown and slot count.

    • allocate

      public static KvCachePool allocate(KvCacheLayout layout, Device device, ScalarType dtype, double memFraction)
      Allocates a pool with DEFAULT_PAGE_SIZE.
      Parameters:
      layout - family-agnostic cache layout.
      device - compute device.
      dtype - cache element dtype.
      memFraction - static-region fraction of total GPU memory.
      Returns:
      the allocated pool.

    • forTesting

      public static KvCachePool forTesting(KvCacheLayout layout, Device device)
      Creates a small pool sized to maxBatchSize × maxSeqLen for unit tests.
      Parameters:
      layout - cache layout.
      device - compute device.
      Returns:
      the test pool.

    • forMtp

      public static KvCachePool forMtp(KvCacheLayout layout, Device device, ScalarType dtype, int pageSize)
      Creates a small fixed-size pool sized to maxBatchSize × maxSeqLen (page-aligned) for the MTP draft head's own short-lived, per-round KV cache. Unlike the main model's allocate(KvCacheLayout, Device, ScalarType, double, int), this is not sized by a fraction of free device memory — the draft window is tiny and fixed — but it must still use the model's real KV cache dtype and page size, not forTesting(KvCacheLayout, Device)'s CPU-unit-test-only defaults (float32, page size 1): a production batched draft step (multiple concurrent requests' MTP rounds sharing one FlashInfer decode call) crashed with an illegal memory access until this matched the main pool's real dtype/page size.
      Parameters:
      layout - cache layout.
      device - compute device.
      dtype - element dtype (the model's real KV cache dtype).
      pageSize - tokens per page (the model's real KV cache page size).
      Returns:
      the MTP pool.

    • forMtpHistory

      public static KvCachePool forMtpHistory(KvCacheLayout layout, Device device, ScalarType dtype, int pageSize, int wantSlots)
      Creates the MTP draft head's persistent, request-bound KV pool (one logical sequence history per bound request, like the main pool but for the draft layer(s) only). Slot count mirrors the main pool's (wantSlots) so any request the main pool can hold fits here, capped to half of the currently free device memory so it can never starve the working set. Prefix reuse is off (the draft history is rebuilt from hidden states, not from radix pages).
      Parameters:
      layout - MTP cache layout (draft layers / local KV heads / head dim).
      device - compute device.
      dtype - element dtype (the model's real KV cache dtype).
      pageSize - tokens per page (the main pool's page size).
      wantSlots - desired slot count (typically the main pool's numSlots()).
      Returns:
      the MTP history pool.
      Throws:
      IllegalStateException - if not even one page fits.

    • radix

      public RadixCache radix()
      Returns the embedded radix tree used for prefix sharing.
      Returns:
      the radix tree.

    • numLayers

      public int numLayers()
      Returns the number of transformer layers covered by this pool.
      Returns:
      layer count.

    • numSlots

      public int numSlots()
      Returns the total number of token slots.
      Returns:
      slot count.

    • freePages

      public int freePages()
      Returns the number of free pages.
      Returns:
      free page count.

    • freeSlots

      public int freeSlots()
      Returns the number of free token slots (freePages × pageSize).
      Returns:
      free slot count.

    • pageSize

      public int pageSize()
      Returns the page size in tokens.
      Returns:
      tokens per page.

    • setDecodeGraphBuffers

      public void setDecodeGraphBuffers(boolean enabled)
      Enables stable-address KV index tensors for batch-1 CUDA graph decode.
      Parameters:
      enabled - true during graph capture / replay forwards.

    • decodeGraphBuffers

      public boolean decodeGraphBuffers()
      Returns whether decode CUDA-graph fixed buffers are active for this step.
      Returns:
      whether decode CUDA-graph fixed buffers are active for this step.

    • setVerifyGraphBuffers

      public void setVerifyGraphBuffers(boolean enabled)
      Enables stable-address KV index/CSR tensors for the MTP window-verify CUDA graph. Independent of setDecodeGraphBuffers(boolean).
      Parameters:
      enabled - true during verify graph capture / replay forwards.

    • setVerifyWindowKernel

      public void setVerifyWindowKernel(boolean enabled)
      Enables or disables eager FlashInfer prefill attention for verify windows.
      Parameters:
      enabled - whether a verify-window forward is in progress.

    • verifyWindowKernel

      public boolean verifyWindowKernel()
      Returns whether a verify-window forward is in progress (see setVerifyWindowKernel(boolean)).
      Returns:
      true during a verify-window forward.

    • verifyGraphBuffers

      public boolean verifyGraphBuffers()
      Returns whether verify CUDA-graph fixed buffers are active for this step.
      Returns:
      whether verify CUDA-graph fixed buffers are active for this step.

    • numPagesForLength

      public int numPagesForLength(int length)
      KV page count for a uniform cache length (CUDA graph bucket key).
      Parameters:
      length - inclusive cached token count.
      Returns:
      number of KV pages spanned by length.

    • stepsUntilPageBoundary

      public int stepsUntilPageBoundary(int cacheLen)
      Decode steps remaining before the next KV page boundary (uniform cache length).
      Parameters:
      cacheLen - inclusive cached token count for the current step.
      Returns:
      0 when the next decode step crosses into a new page bucket.

    • firstCacheLenInNextPageBucket

      public int firstCacheLenInNextPageBucket(int numPages)
      First inclusive cache length in the page bucket after numPages.
      Parameters:
      numPages - current KV page count.
      Returns:
      cache length for the first step in bucket numPages + 1.

    • beginPrefetchDecodeGraphStep

      public void beginPrefetchDecodeGraphStep(int cacheLen, int cachePos, int batch)
      Uses live request page-table layout for attention and scratch slots for KV writes so prefetch capture sees the same CSR tensors replay will use.
      Parameters:
      cacheLen - inclusive cached token count for the prefetch step.
      cachePos - write position within the cache length.
      batch - decode batch size.

    • currentStepFlashInferMetadata

      public FlashInferKvMetadata currentStepFlashInferMetadata()
      Returns the shared step FlashInfer metadata.
      Returns:
      shared step metadata (do not close).

    • installPrefetchedStepMetadata

      public void installPrefetchedStepMetadata(FlashInferKvMetadata meta, int uniformLen)
      Installs CSR metadata captured during prefetch so graph replay reuses the same tensor addresses as capture.
      Parameters:
      meta - prefetched FlashInfer CSR metadata.
      uniformLen - uniform inclusive cache length for the step.

    • endPrefetchDecodeGraphStep

      public void endPrefetchDecodeGraphStep()
      Restores live step metadata after beginPrefetchDecodeGraphStep(int, int, int).

    • prepareDecodeGraphStep

      public void prepareDecodeGraphStep(int cacheLen, int startPos)
      Updates FlashInfer metadata and the static KV index before graph replay.
      Parameters:
      cacheLen - inclusive cache length for the active request(s).
      startPos - KV write position for this decode step.

    • prepareDecodeGraphStep

      public void prepareDecodeGraphStep(int cacheLen, int[] cachePositions)
      Updates FlashInfer metadata and per-row KV indices before graph replay.
      Parameters:
      cacheLen - inclusive cache length (uniform across the batch).
      cachePositions - KV write position per batch row.

    • decodeGraphMetadata

      public FlashInferKvMetadata decodeGraphMetadata(int batch)
      The fixed-address page table decode graphs read (see prepareDecodeGraphStep(int, int)).
      Parameters:
      batch - decode batch size.
      Returns:
      the metadata; its tensors keep their addresses for the life of the pool.
      Throws:
      IllegalStateException - if prepareDecodeGraphStep(int, int) has not run for this batch size.

    • prepareVerifyGraphStep

      public void prepareVerifyGraphStep(int cacheLen, int startPos, int windowLen)
      Updates FlashInfer metadata and the flat [batch*windowLen] KV write-index buffer before a verify-graph replay (single-request form).
      Parameters:
      cacheLen - inclusive cache length for the active request.
      startPos - KV write position for the first token of the window.
      windowLen - number of tokens written by this verify step.

    • prepareVerifyGraphStep

      public void prepareVerifyGraphStep(int cacheLen, int[] startPositions, int windowLen)
      Updates FlashInfer metadata and the flat [batch*windowLen] KV write-index buffer before a verify-graph replay.
      Parameters:
      cacheLen - inclusive cache length (uniform across the batch).
      startPositions - KV write position of the window's first token per batch row.
      windowLen - number of tokens written by this verify step (uniform across the batch).

    • prepareVerifyGraphStepRagged

      public void prepareVerifyGraphStepRagged(int[] cacheLengths, int[] startPositions, int windowLen)
      Ragged form of prepareVerifyGraphStep(int, int[], int): rows (concurrent requests) may be at different cache lengths. Rewrites the graph's fixed-address CSR and KV write index.
      Parameters:
      cacheLengths - inclusive cache length per batch row.
      startPositions - KV write position of each row's first window token.
      windowLen - tokens written per row.

    • verifyGraphMetadata

      public FlashInferKvMetadata verifyGraphMetadata()
      Returns the verify CUDA graph's fixed-address CSR (valid after prepareVerifyGraphStep(int, int[], int)); do not close.
      Returns:
      persistent verify-graph CSR metadata.

    • consumeVerifyGraphMetaRealloc

      public boolean consumeVerifyGraphMetaRealloc()
      Returns and clears whether verifyGraphMetadata() was reallocated (new addresses) since the last call; a captured graph referencing the old tensors must then be discarded.
      Returns:
      true when the persistent CSR tensors were replaced.

    • prepareVerifyGraphStep

      public void prepareVerifyGraphStep(int[] cacheLengths, int[] startPositions, int windowLen)
      Updates FlashInfer metadata and the flat [batch*windowLen] KV write-index buffer before a batched eager verify forward whose rows belong to different concurrent requests at different absolute cache lengths — the multi-request counterpart of prepareVerifyGraphStep(int, int[], int), which requires a single cacheLen shared by every row (fine for the single-request / graph-capturable cohort that method still serves).
      Parameters:
      cacheLengths - inclusive cache length per batch row.
      startPositions - KV write position of the window's first token per batch row.
      windowLen - number of tokens written by this verify step (uniform across the batch).

    • verifyQoIndptrBuf

      public Tensor verifyQoIndptrBuf(int batch, int windowLen)
      Returns a stable query-side CSR [batch+1] tensor with values [0, windowLen, 2*windowLen, ...], rebuilt (new address) only when the (batch, windowLen) bucket changes. Callers that need a fixed address across a captured verify graph's lifetime must only call this once per bucket, before capture, and reuse the returned tensor for every replay in that bucket.
      Parameters:
      batch - active batch size.
      windowLen - number of query tokens per row.
      Returns:
      query-side CSR tensor; do not close (owned by the pool).

    • requestCapacity

      public int requestCapacity()
      Capacity reserved for the currently bound request, or 0 if none. Generation must not access positions >= this value.
      Returns:
      bound request capacity in tokens.

    • setPrefixReuseEnabled

      public void setPrefixReuseEnabled(boolean enabled)
      Enables or disables SGLang-style prefix reuse via bindWithPrefix(int[], int). When disabled, bindWithPrefix(int[], int) behaves like contiguous bindRequests(int, int)(1, totalCapacity).
      Parameters:
      enabled - true to match/insert against the radix tree.

    • isPrefixReuseEnabled

      public boolean isPrefixReuseEnabled()
      Returns whether prefix reuse is enabled.
      Returns:
      true if prefix reuse is enabled.

    • prefixPromptTokens

      public long prefixPromptTokens()
      Cumulative prompt tokens seen by prefix bind (full length, not page-floored).
      Returns:
      cumulative prompt token count.

    • prefixMatchTokens

      public long prefixMatchTokens()
      Cumulative matched prefix tokens.
      Returns:
      cumulative matched token count.

    • prefixInsertTokens

      public long prefixInsertTokens()
      Cumulative tokens inserted into the radix tree.
      Returns:
      cumulative inserted token count.

    • bindRequests

      public void bindRequests(int batchSize, int capacity)
      Reserves a page-aligned slot map for each item in a batch. Must be called before put(int, int, Tensor, Tensor)/get(int, int) for a request. Previously bound slots (including multi-request bindings) are released without radix insert. Physical pages need not be contiguous (paged attention).

      Capacity is page-aligned and clamped to free pages in the static pool (after radix eviction). The pool never grows. Callers should read requestCapacity() and stop generation at that limit.

      Parameters:
      batchSize - number of parallel requests.
      capacity - desired slots reserved per request (prompt + generation).
      Throws:
      IllegalStateException - if no page can be reserved per request.

    • bindWithPrefix

      public int bindWithPrefix(int[] promptTokens, int totalCapacity)
      Match a prompt against the radix tree, lock the matched node, allocate suffix/decode slots, and bind a slot map for batch size 1.

      Clears any multi-request bindings first so LanguageModel.generate(int[], int, double, double, boolean, long, GenerationListener) remains exclusive.

      Desired totalCapacity is clamped to matched prefix plus free private pages. The static pool never grows. If the prompt itself cannot fit after clamping, this method throws.

      Parameters:
      promptTokens - prompt token ids.
      totalCapacity - desired slots (prompt + generation headroom).
      Returns:
      page-aligned matched prefix length (0 if miss or reuse disabled).
      Throws:
      IllegalArgumentException - if the prompt cannot fit in free KV slots.

    • bindRequest

      public int bindRequest(int[] promptTokens, int totalCapacity)
      Allocates slots for one request without clearing other multi-request bindings. If the multi-request map was empty, activates this request.

      Requires the full totalCapacity (page-aligned). When free KV is insufficient, throws KvCacheExhaustedException so the caller can leave the job queued until Instant Eviction frees pages.

      Parameters:
      promptTokens - prompt token ids.
      totalCapacity - desired slots (prompt + generation headroom).
      Returns:
      request id (> 0).
      Throws:
      KvCacheExhaustedException - if the pool cannot reserve totalCapacity.

    • isBound

      public boolean isBound(int requestId)
      Returns whether requestId currently has a multi-request binding.
      Parameters:
      requestId - id returned by bindRequest(int[], int).
      Returns:
      true when bound.

    • matchedPrefixLen

      public int matchedPrefixLen(int requestId)
      Page-aligned matched prefix length for a multi-request binding.
      Parameters:
      requestId - id returned by bindRequest(int[], int).
      Returns:
      matched prefix length, or 0 if unknown.

    • unbindRequest

      public void unbindRequest(int requestId)
      Instant Eviction: free this request's private pages; unlock its radix node.
      Parameters:
      requestId - id returned by bindRequest(int[], int).

    • finishRequest

      public void finishRequest(int requestId, int[] sequenceTokens)
      Inserts the sequence into the radix tree (prefix reuse) then unbindRequest(int).
      Parameters:
      requestId - id returned by bindRequest(int[], int).
      sequenceTokens - prompt and generated token ids (no pad).

    • activateStep

      public void activateStep(int... requestIds)
      Sets the active step batch used by put(int, int, Tensor, Tensor)/get(int, int)/ buildFlashInferMetadata(int)/ensureRequest(int, int) to the given request ids (in order). Capacities may differ per request; ensureRequest(int, int) verifies each requestSlots[b].length >= endPos.
      Parameters:
      requestIds - bound request ids from bindRequest(int[], int).

    • boundRequestCount

      public int boundRequestCount()
      Number of currently bound requests (multi-request map size, or legacy exclusive batch size when the map is empty).
      Returns:
      bound request count.

    • finishRequest

      public void finishRequest(int[] sequenceTokens)
      Inserts prompt+completion into the radix tree, frees duplicate/private pages not retained by the tree, and unlocks the matched node.

      If exactly one multi-request binding exists, finishes that binding. Otherwise finishes the sole active legacy binding (LanguageModel.generate).

      Parameters:
      sequenceTokens - prompt and generated token ids (no pad).

    • unbindRequests

      public void unbindRequests()
      Releases slots reserved by bindRequests(int, int) / bindWithPrefix(int[], int) (and any multi-request bindings) without inserting into the radix tree. Tree-owned matched pages are not freed.

    • ensureRequest

      public void ensureRequest(int batchSize, int endPos)
      Ensures a request binding exists and is large enough for endPos. Never grows the static pool: if capacity is insufficient, throws KvCacheExhaustedException. Callers must bind up front and stop generation at requestCapacity().

      When capacities differ across the active step, each requestSlots[b].length must be >= endPos.

      Parameters:
      batchSize - number of parallel requests.
      endPos - exclusive end position that will be accessed.
      Throws:
      KvCacheExhaustedException - if unbound or endPos exceeds capacity.

    • put

      public void put(int layer, int startPos, Tensor k, Tensor v)
      Writes key/value activations for the bound batch at positions [startPos, startPos + seqlen).
      Parameters:
      layer - layer index.
      startPos - starting position within each request's reserved range.
      k - keys shaped [batch, seqlen, numKvHeads, headDim].
      v - values shaped [batch, seqlen, numKvHeads, headDim].

    • put

      public void put(int layer, int[] startPos, Tensor k, Tensor v)
      Writes key/value activations at a per-row start position (typically decode seqlen == 1 with unequal positions across the batch).
      Parameters:
      layer - layer index.
      startPos - write start position per batch row (length == batch).
      k - keys shaped [batch, seqlen, numKvHeads, headDim].
      v - values shaped [batch, seqlen, numKvHeads, headDim].

    • put

      public void put(int layer, int[] startPos, Tensor k, Tensor v, Tensor slotIdx)
      Like put(int, int[], Tensor, Tensor) with the physical slot of every written position already on the device ([batch * seqlen], row-major), so a model can upload it once per step instead of once per layer. Ignored when a decode/verify CUDA-graph index buffer is active.
      Parameters:
      layer - layer index.
      startPos - write start position per batch row (length == batch).
      k - keys shaped [batch, seqlen, numKvHeads, headDim].
      v - values shaped [batch, seqlen, numKvHeads, headDim].
      slotIdx - device tensor of slots, or null to build it from the host.

    • slotIndices

      public long[] slotIndices(int length)
      Returns the physical slot of every position [0, length) of the bound batch, row-major [batch * length]. Side caches (e.g. MSA index keys) use these slots so they share allocation, prefix reuse and eviction with K/V. Rows whose reserved capacity is shorter than length (requests of different lengths decoded in one batch) repeat their last slot for the extra positions; the caller must mask those positions.
      Parameters:
      length - number of positions per row.
      Returns:
      slot ids.

    • slotIndices

      public long[] slotIndices(int[] startPos, int seqlen)
      Returns the physical slot of positions [startPos[b], startPos[b] + seqlen) for each bound row, row-major [batch * seqlen], matching the layout of put(int, int[], Tensor, Tensor). Slots are reserved lazily by put, so call this after the K/V write for the same window.
      Parameters:
      startPos - write start position per row.
      seqlen - positions per row.
      Returns:
      slot ids.

    • get

      public Tuple2<Tensor,Tensor> get(int layer, int length)
      Reads key/value activations for positions [0, length) of the bound batch.
      Parameters:
      layer - layer index.
      length - number of positions to read.
      Returns:
      (keys, values) each shaped [batch, length, numKvHeads, headDim].

    • keyLayer

      public Tensor keyLayer(int layer)
      One layer of the key cache [numSlots, numKvHeads, headDim] in the pool dtype, as a view.
      Parameters:
      layer - layer index.
      Returns:
      the view; the caller closes it.

    • valueLayer

      public Tensor valueLayer(int layer)
      One layer of the value cache; see keyLayer(int).
      Parameters:
      layer - layer index.
      Returns:
      the view; the caller closes it.

    • kvIsComputeDtype

      public boolean kvIsComputeDtype()
      Returns whether K/V are stored in the compute dtype (not quantized FP8), so kernels can read them.
      Returns:
      true when K/V are stored in the compute dtype (not quantized FP8), so kernels can read them.

    • getPadded

      public Tuple2<Tensor,Tensor> getPadded(int layer, int length)
      Like get(int, int) for a batch whose rows have different reserved capacities: positions past a row's own capacity read that row's last slot instead of failing. The caller must mask such positions (a decode step masks everything beyond each row's current position).
      Parameters:
      layer - layer index.
      length - number of positions to read per row.
      Returns:
      (keys, values) each shaped [batch, length, numKvHeads, headDim].

    • getPadded

      public Tuple2<Tensor,Tensor> getPadded(int layer, int length, Tensor idx)
      getPadded(int, int) with the [batch * length] slot index already on the device.
      Parameters:
      layer - layer index.
      length - number of positions to read per row.
      idx - device tensor from slotIndices(int); not closed here.
      Returns:
      (keys, values) each shaped [batch, length, numKvHeads, headDim].

    • clearStepFlashInferMetadata

      public void clearStepFlashInferMetadata()
      Releases step-scoped FlashInfer metadata (if any). Called automatically at the start of each activateStep(int...).

    • sealLength

      public boolean sealLength(int sealedLen)
      Seals the readable KV length for the active step without freeing pages. Rebuilds FlashInfer CSR for sealedLen so attention cannot see positions >= sealedLen (required after speculative reject).
      Parameters:
      sealedLen - inclusive committed token count ([0, sealedLen)).
      Returns:
      true when the existing CSR was bumped in place (same tensor addresses survive); false when it was rebuilt (a genuine page-boundary crossing, or a ragged step) — callers holding a captured CUDA graph keyed to the old CSR addresses must invalidate it when this returns false.

    • invalidateRange

      public void invalidateRange(int fromPos, int endPos)
      Zeros K/V for positions [fromPos, endPos) on every layer of the active batch (defense-in-depth for rejected speculative tails that share a KV page with committed tokens).
      Parameters:
      fromPos - inclusive start position.
      endPos - exclusive end position.

    • zeroRejectedKv

      public static boolean zeroRejectedKv()
      Whether rejected speculative tails are zeroed (see ZERO_REJECTED_KV).
      Returns:
      true when zeroing is enabled.

    • truncateTo

      public boolean truncateTo(int sealedLen, int writtenEnd)
      After a speculative window write of length writtenEnd, seal attention to sealedLen and zero the rejected tail [sealedLen, writtenEnd). Does not free pages or shrink capacity.
      Parameters:
      sealedLen - committed inclusive length after accept/reject.
      writtenEnd - exclusive end of the speculative window write.
      Returns:
      true when the existing CSR was bumped in place (see sealLength(int)); false when it was rebuilt.

    • sharedFlashInferMetadata

      public FlashInferKvMetadata sharedFlashInferMetadata(int length)
      Returns shared FlashInfer metadata for a uniform cache length within the current step. Do not close the returned value.
      Parameters:
      length - number of cached positions (inclusive of the current step).
      Returns:
      CSR metadata reused for all attention layers in this step.

    • sharedFlashInferMetadata

      public FlashInferKvMetadata sharedFlashInferMetadata(int[] lengths)
      Returns shared FlashInfer metadata for ragged per-row cache lengths within the current step. Do not close the returned value.
      Parameters:
      lengths - cached length per active request.
      Returns:
      CSR metadata reused for all attention layers in this step.

    • buildFlashInferMetadata

      public FlashInferKvMetadata buildFlashInferMetadata(int length)
      Builds FlashInfer CSR metadata for a uniform cache length.
      Parameters:
      length - inclusive cached token count.
      Returns:
      FlashInfer KV metadata for the bound batch.

    • buildFlashInferMetadata

      public FlashInferKvMetadata buildFlashInferMetadata(int[] lengths)
      Builds FlashInfer CSR page-table metadata with a per-row cached length (ragged decode batches).
      Parameters:
      lengths - cached length per active request (length[b] == cacheLen).
      Returns:
      CSR metadata; caller owns and must close.

    • keyCache

      public Tensor keyCache()
      Returns the key cache buffer [numLayers, numSlots, numKvHeads, headDim].
      Returns:
      key cache.

    • valueCache

      public Tensor valueCache()
      Returns the value cache buffer [numLayers, numSlots, numKvHeads, headDim].
      Returns:
      value cache.

    • numKvHeads

      public int numKvHeads()
      Returns the number of KV heads stored in the pool.
      Returns:
      number of KV heads stored in the pool.

    • headDim

      public int headDim()
      Returns the head dimension.
      Returns:
      head dimension.

    • device

      public Device device()
      Returns the hosting device.
      Returns:
      hosting device.

    • flashInferWorkspace

      public FlashInferWorkspace flashInferWorkspace()
      Returns (and lazily creates) the FlashInfer workspace for this pool.
      Returns:
      workspace, or null when FlashInfer is unavailable.

    • alloc

      public long[] alloc(int numTokens)
      Allocates numTokens slots (page-aligned) and returns their indices. Used by the inference engine when inserting into the radix tree.

      Pages are packed from the free list and need not be physically contiguous. Within each page, slot indices remain pageId * pageSize + offset so FlashInfer / gather paths can treat the result as a standard page table.

      Parameters:
      numTokens - number of tokens to allocate.
      Returns:
      slot indices of length alignedLen.

    • free

      public void free(long[] slots)
      Returns pages covering the given slot indices to the free list. Slot indices should be page-aligned groups.
      Parameters:
      slots - slot indices to free.

    • close

      public void close()
      Specified by:
      close in interface AutoCloseable