Class DeltaNetStatePool

java.lang.Object
smile.llm.model.qwen.DeltaNetStatePool
All Implemented Interfaces:
AutoCloseable
public class DeltaNetStatePool extends Object implements AutoCloseable
Per-request DeltaNet recurrent and causal-conv state for linear-attention layers.

Recurrent state is typically float32 (fused CUDA kernel / stable in-place updates). Conv left-context must match the compute dtype (bf16/fp16) so decode concat(convState, hidden) does not promote activations to float and break later bf16 linear layers.

Multi-request continuous batching assigns each requestId a stable home row. activateStep(int...) packs those rows into [0, B) for the mixer; scatterActive() writes them back after the forward.

  • Constructor Summary

    Constructors
    Constructor
    Description
    DeltaNetStatePool(int numLinearLayers, int numVHeads, int keyHeadDim, int valueHeadDim, int convDim, int convKernel, int maxBatchSize, Device device, ScalarType dtype)
    Constructor using the same dtype for recurrent and conv buffers (tests).
    DeltaNetStatePool(int numLinearLayers, int numVHeads, int keyHeadDim, int valueHeadDim, int convDim, int convKernel, int maxBatchSize, Device device, ScalarType recurrentDtype, ScalarType convDtype)
    Constructor with separate recurrent / conv dtypes.

  • Method Summary

    Modifier and Type
    Method
    Description
    void
    activateStep(int... requestIds)
    Packs bound request rows into working slots [0, B) for a forward.
    Tensor
    activeConv(int linearLayerId)
    Conv rows packed by activateStep(int...) into [0, boundBatch).
    Tensor
    activeRecurrent(int linearLayerId)
    Recurrent rows packed by activateStep(int...) into [0, boundBatch).
    int
    bindRequest(int requestId)
    Binds a stable home row for requestId and zeros it.
    int
    boundBatch()
    Returns the bound batch size.
    int
    boundRequestCount()
    Returns the number of multi-request bindings.
    void
    close()
     
    Tensor
    conv(int linearLayerId)
    Returns the conv state buffer for a linear-attention layer.
    boolean
    ensureSpeculativeCheckpointRange(int firstSlot, int lastSlot)
    Allocates checkpoint slots firstSlot..lastSlot (inclusive) only, for max(1, boundBatch) rows.
    boolean
    ensureSpeculativeCheckpointRange(int firstSlot, int lastSlot, int rows)
    Allocates checkpoint slots firstSlot..lastSlot for an explicit number of rows, so the memory can be reserved at startup (before the KV pool is sized) instead of being grown on demand once the GPU is nearly full.
    boolean
    ensureSpeculativeCheckpoints(int numSlots)
    Ensures numSlots mid-verify checkpoint buffers sized for the currently activated batch (boundBatch), not the full pool maxBatchSize.
    boolean
    hasSpeculativeCheckpoints(int slots, int rows)
    Whether per-step checkpoints for slots slots and rows rows are allocated.
    boolean
    hasSpeculativeRange(int firstSlot, int lastSlot, int rows)
    Whether checkpoint slots firstSlot..lastSlot are stored for at least rows rows.
    int
    numLinearLayers()
    Returns the linear-attention layer count.
    Tensor
    recurrent(int linearLayerId)
    Returns the recurrent state buffer for a linear-attention layer.
    void
    releaseSpeculativeCheckpoints()
    Releases speculative checkpoint buffers (call after a verify round when GPU headroom is tight).
    void
    reset(int batchSize)
    Zeros all states and records the active batch size (legacy exclusive generate).
    void
    restoreCheckpoint(int slot)
    Restores active working rows from checkpoint slot.
    void
    restoreCheckpointPerRow(int[] slots)
    Restores each active row i from its own checkpoint slot slots[i] — the batched-cohort counterpart of restoreCheckpoint(int), which restores every active row from the same slot.
    void
    saveCheckpoint(int slot)
    Copies active working rows [0, boundBatch) into checkpoint slot.
    void
    saveCheckpointForLayer(int slot, int layerId)
    Copies active working rows [0, boundBatch) for a single linear-attention layer into checkpoint slot.
    void
    scatterActive()
    Writes working slots [0, B) back to each request's home row.
    void
    setVerifyWindowActive(boolean active)
    Sets whether the primary MTP verify-window forward is in flight.
    long
    speculativeCheckpointBytes(int numSlots, int rows)
    Device bytes of the speculative checkpoint buffers for numSlots slots of rows rows each (all linear layers, recurrent plus conv state).
    long
    speculativeCheckpointGrowthBytes(int numSlots, int rows)
    Additional device bytes ensureSpeculativeCheckpoints(int) would have to allocate for numSlots x max(1, rows) (0 when the existing buffers already suffice).
    long
    speculativeCheckpointGrowthBytes(int firstSlot, int lastSlot, int rows)
    Additional device bytes ensureSpeculativeCheckpointRange(int, int) would have to allocate for slots firstSlot..lastSlot of max(1, rows) rows (0 when already sufficient).
    Tensor
    speculativeConvSlot(int slot, int layerId)
    Returns checkpoint slot's full conv tensor for a linear layer ([rows, C, K-1]), or null when absent.
    Tensor
    speculativeRecurrentSlot(int slot, int layerId)
    Returns checkpoint slot's full recurrent tensor for a linear layer ([rows, H, K, V], float, contiguous), or null when absent.
    void
    unbind()
    Clears the active-request binding after exclusive generate finishes.
    void
    unbindRequest(int requestId)
    Releases the home row for requestId.
    boolean
    verifyWindowActive()
    Returns whether the primary MTP verify-window forward is in flight.
    void
    withPreservedActive(Runnable action)
    Runs action without permanently mutating the active DeltaNet working rows.

    Methods inherited from class Object

    clone, equals, finalize, getClass, hashCode, notify, notifyAll, toString, wait, wait, wait

  • Constructor Details

    • DeltaNetStatePool

      public DeltaNetStatePool(int numLinearLayers, int numVHeads, int keyHeadDim, int valueHeadDim, int convDim, int convKernel, int maxBatchSize, Device device, ScalarType dtype)
      Constructor using the same dtype for recurrent and conv buffers (tests).
      Parameters:
      numLinearLayers - number of linear-attention layers.
      numVHeads - number of value heads.
      keyHeadDim - key head dimension.
      valueHeadDim - value head dimension.
      convDim - causal conv channel dimension.
      convKernel - causal conv kernel size.
      maxBatchSize - maximum concurrent home rows.
      device - storage device.
      dtype - dtype for both recurrent and conv buffers.

    • DeltaNetStatePool

      public DeltaNetStatePool(int numLinearLayers, int numVHeads, int keyHeadDim, int valueHeadDim, int convDim, int convKernel, int maxBatchSize, Device device, ScalarType recurrentDtype, ScalarType convDtype)
      Constructor with separate recurrent / conv dtypes.
      Parameters:
      numLinearLayers - number of linear-attention layers.
      numVHeads - number of value heads.
      keyHeadDim - key head dimension.
      valueHeadDim - value head dimension.
      convDim - causal conv channel dimension.
      convKernel - causal conv kernel size.
      maxBatchSize - maximum concurrent home rows.
      device - storage device.
      recurrentDtype - dtype for recurrent state buffers.
      convDtype - dtype for conv state buffers.

  • Method Details

    • reset

      public void reset(int batchSize)
      Zeros all states and records the active batch size (legacy exclusive generate).
      Parameters:
      batchSize - active batch size.

    • bindRequest

      public int bindRequest(int requestId)
      Binds a stable home row for requestId and zeros it.
      Parameters:
      requestId - id aligned with KvCachePool.bindRequest(int[], int).
      Returns:
      home row index.

    • unbindRequest

      public void unbindRequest(int requestId)
      Releases the home row for requestId.
      Parameters:
      requestId - previously bound id.

    • activateStep

      public void activateStep(int... requestIds)
      Packs bound request rows into working slots [0, B) for a forward.

      Working slots are the same storage as the first B home rows, so packing has to be careful in two ways. First, all sources are read before any is written (a sequential home-to-slot copy corrupts the cohort itself when homes are not in ascending order, e.g. once freed rows have been reused by later requests). Second, any live request outside the cohort whose home row lies in [0, B) is stashed and put back by scatterActive(), so a forward over one cohort never disturbs another's state (a speculating group and a plain-decode group alternate within one engine tick, for example). When the cohort already occupies rows 0..B-1 in order nothing is copied.

      Parameters:
      requestIds - bound request ids (order = batch).

    • setVerifyWindowActive

      public void setVerifyWindowActive(boolean active)
      Sets whether the primary MTP verify-window forward is in flight.
      Parameters:
      active - true around the verify-window forward only.

    • verifyWindowActive

      public boolean verifyWindowActive()
      Returns whether the primary MTP verify-window forward is in flight.
      Returns:
      true when GatedDeltaNet should run its per-position verify loop instead of the batched S>1 path.

    • scatterActive

      public void scatterActive()
      Writes working slots [0, B) back to each request's home row. Call after every forward that used activateStep(int...).

    • withPreservedActive

      public void withPreservedActive(Runnable action)
      Runs action without permanently mutating the active DeltaNet working rows.

      Used for CUDA graph prefetch forwards that share the current activateStep(int...) packing.

      Parameters:
      action - work to run while active rows are preserved.

    • speculativeCheckpointBytes

      public long speculativeCheckpointBytes(int numSlots, int rows)
      Device bytes of the speculative checkpoint buffers for numSlots slots of rows rows each (all linear layers, recurrent plus conv state).
      Parameters:
      numSlots - checkpoint slots.
      rows - batch rows per slot.
      Returns:
      bytes the buffers occupy on this pool's device.

    • speculativeCheckpointGrowthBytes

      public long speculativeCheckpointGrowthBytes(int numSlots, int rows)
      Additional device bytes ensureSpeculativeCheckpoints(int) would have to allocate for numSlots x max(1, rows) (0 when the existing buffers already suffice).
      Parameters:
      numSlots - checkpoint slots required.
      rows - batch rows required.
      Returns:
      extra bytes needed.

    • speculativeCheckpointGrowthBytes

      public long speculativeCheckpointGrowthBytes(int firstSlot, int lastSlot, int rows)
      Additional device bytes ensureSpeculativeCheckpointRange(int, int) would have to allocate for slots firstSlot..lastSlot of max(1, rows) rows (0 when already sufficient).
      Parameters:
      firstSlot - first slot required.
      lastSlot - last slot required.
      rows - batch rows required.
      Returns:
      extra bytes needed.

    • ensureSpeculativeCheckpoints

      public boolean ensureSpeculativeCheckpoints(int numSlots)
      Ensures numSlots mid-verify checkpoint buffers sized for the currently activated batch (boundBatch), not the full pool maxBatchSize. Full-pool sizing OOMs under serving configs.

      Growth-only: never reallocates smaller for a later, smaller boundBatch (the early-return below already guarantees this), since a shrink would be pointless — the real reason growth matters is that any verify CUDA graph already captured against the old, smaller buffers becomes invalid the moment this reallocates (replay would then write into freed memory) — see the true-return contract below.

      Parameters:
      numSlots - checkpoint slots (N+1 for N draft tokens).
      Returns:
      true if this call reallocated (the caller must then invalidate any captured verify CUDA graph via QwenModel.invalidateVerifyCudaGraphs() before it can be replayed again).

    • ensureSpeculativeCheckpointRange

      public boolean ensureSpeculativeCheckpointRange(int firstSlot, int lastSlot)
      Allocates checkpoint slots firstSlot..lastSlot (inclusive) only, for max(1, boundBatch) rows. Slot numbers keep their meaning ("state after window position slot-1", slot 0 = pre-window), but slots outside the range have no storage: a save to one is a no-op and a restore from one is an error. In checkpoint-replay mode only slots 1..n are ever restored from (slot 0 is never read, and the last slot equals the working state), so allocating just those cuts the memory by 2/(n+2).
      Parameters:
      firstSlot - first slot to store.
      lastSlot - last slot to store.
      Returns:
      true if this call reallocated (see ensureSpeculativeCheckpoints(int)).

    • ensureSpeculativeCheckpointRange

      public boolean ensureSpeculativeCheckpointRange(int firstSlot, int lastSlot, int rows)
      Allocates checkpoint slots firstSlot..lastSlot for an explicit number of rows, so the memory can be reserved at startup (before the KV pool is sized) instead of being grown on demand once the GPU is nearly full.
      Parameters:
      firstSlot - first slot to store.
      lastSlot - last slot to store.
      rows - batch rows to provision.
      Returns:
      true if this call reallocated.

    • hasSpeculativeRange

      public boolean hasSpeculativeRange(int firstSlot, int lastSlot, int rows)
      Whether checkpoint slots firstSlot..lastSlot are stored for at least rows rows.
      Parameters:
      firstSlot - first checkpoint slot of the range (inclusive).
      lastSlot - last checkpoint slot of the range (inclusive).
      rows - minimum number of rows the checkpoints must cover.
      Returns:
      true if the range is fully stored with at least rows rows.

    • releaseSpeculativeCheckpoints

      public void releaseSpeculativeCheckpoints()
      Releases speculative checkpoint buffers (call after a verify round when GPU headroom is tight).

    • saveCheckpoint

      public void saveCheckpoint(int slot)
      Copies active working rows [0, boundBatch) into checkpoint slot.
      Parameters:
      slot - checkpoint index (0 .. numSlots-1).

    • restoreCheckpoint

      public void restoreCheckpoint(int slot)
      Restores active working rows from checkpoint slot.
      Parameters:
      slot - checkpoint index (0 .. numSlots-1).

    • restoreCheckpointPerRow

      public void restoreCheckpointPerRow(int[] slots)
      Restores each active row i from its own checkpoint slot slots[i] — the batched-cohort counterpart of restoreCheckpoint(int), which restores every active row from the same slot. Needed when concurrent requests verified together in one round accept different numbers of draft tokens.
      Parameters:
      slots - checkpoint slot per active row (length boundBatch).

    • speculativeRecurrentSlot

      public Tensor speculativeRecurrentSlot(int slot, int layerId)
      Returns checkpoint slot's full recurrent tensor for a linear layer ([rows, H, K, V], float, contiguous), or null when absent.
      Parameters:
      slot - checkpoint index.
      layerId - ordinal among linear-attention layers.
      Returns:
      the checkpoint tensor (owned by the pool; do not close), or null.

    • speculativeConvSlot

      public Tensor speculativeConvSlot(int slot, int layerId)
      Returns checkpoint slot's full conv tensor for a linear layer ([rows, C, K-1]), or null when absent.
      Parameters:
      slot - checkpoint index.
      layerId - ordinal among linear-attention layers.
      Returns:
      the checkpoint tensor (owned by the pool; do not close), or null.

    • hasSpeculativeCheckpoints

      public boolean hasSpeculativeCheckpoints(int slots, int rows)
      Whether per-step checkpoints for slots slots and rows rows are allocated.
      Parameters:
      slots - number of checkpoint slots to check.
      rows - minimum number of rows the checkpoints must cover.
      Returns:
      true if checkpoints for slots slots and rows rows are allocated.

    • saveCheckpointForLayer

      public void saveCheckpointForLayer(int slot, int layerId)
      Copies active working rows [0, boundBatch) for a single linear-attention layer into checkpoint slot.

      Used by the verify-window per-position loop: QwenModel's layer stack is depth-sequential, so each layer's own forward (including its internal per-position loop) finishes before the next layer starts. A whole-pool save mid-loop would capture sibling layers' stale state; this saves only the layer that just finished its own position t.

      Parameters:
      slot - checkpoint index (0 .. numSlots-1).
      layerId - ordinal among linear-attention layers.

    • unbind

      public void unbind()
      Clears the active-request binding after exclusive generate finishes.

    • boundBatch

      public int boundBatch()
      Returns the bound batch size.
      Returns:
      bound batch size, or 0 if unbound.

    • boundRequestCount

      public int boundRequestCount()
      Returns the number of multi-request bindings.
      Returns:
      number of multi-request bindings.

    • recurrent

      public Tensor recurrent(int linearLayerId)
      Returns the recurrent state buffer for a linear-attention layer.
      Parameters:
      linearLayerId - ordinal among linear-attention layers.
      Returns:
      recurrent state [maxBatch, V, Kdim, Vdim] (first boundBatch rows active).

    • activeRecurrent

      public Tensor activeRecurrent(int linearLayerId)
      Recurrent rows packed by activateStep(int...) into [0, boundBatch).

      Mixer forwards must use this (not recurrent) so batch matmul sees state.shape()[0] == query.shape()[0].

      Parameters:
      linearLayerId - ordinal among linear-attention layers.
      Returns:
      view [boundBatch, V, Kdim, Vdim] into the pool buffer.

    • activeConv

      public Tensor activeConv(int linearLayerId)
      Conv rows packed by activateStep(int...) into [0, boundBatch).
      Parameters:
      linearLayerId - ordinal among linear-attention layers.
      Returns:
      view [boundBatch, C, K-1], or null if unused.

    • conv

      public Tensor conv(int linearLayerId)
      Returns the conv state buffer for a linear-attention layer.
      Parameters:
      linearLayerId - ordinal among linear-attention layers.
      Returns:
      conv state [maxBatch, C, K-1], or null if unused.

    • numLinearLayers

      public int numLinearLayers()
      Returns the linear-attention layer count.
      Returns:
      linear layer count.

    • close

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