Class MtpAnchorPool

java.lang.Object
smile.llm.model.qwen.MtpAnchorPool
All Implemented Interfaces:
AutoCloseable
public class MtpAnchorPool extends Object implements AutoCloseable
Per-request MTP draft anchor (post-final-norm hidden feeding the MTP head's first draft step), keyed by request id.

QwenModel.lastPreNormHidden is a single field shared by every forward (decode/prefill/verify), for whichever batch of requests that forward just processed — safe only because at most one request was ever mid-speculative-round at a time. Real multi-request batching needs each request's anchor kept durably in its own row, scattered here immediately after the forward that produced it (see Qwen.scatterMtpAnchor), so a later forward for a different batch can never clobber it.

  • Constructor Summary

    Constructors
    Constructor
    Description
    MtpAnchorPool(int maxBatchSize, int dim)
    Constructor.

  • Method Summary

    Modifier and Type
    Method
    Description
    int
    bindRequest(int requestId)
    Binds a stable home row for requestId.
    void
    close()
     
    Tensor
    getRow(int requestId)
    Returns a fresh [D] copy of requestId's anchor row, or null if never written.
    Tensor
    getRows(int[] requestIds)
    Returns a fresh [B, D] tensor gathering each request's anchor row in order (batched MTP draft input), or null if never written.
    boolean
    hasRow(int requestId)
    Returns whether requestId has a written anchor row yet.
    void
    setRow(int requestId, Tensor row)
    Writes row's content into requestId's durable anchor.
    void
    unbindRequest(int requestId)
    Releases the home row for requestId.

    Methods inherited from class Object

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

  • Constructor Details

    • MtpAnchorPool

      public MtpAnchorPool(int maxBatchSize, int dim)
      Constructor.
      Parameters:
      maxBatchSize - maximum concurrently bound requests.
      dim - hidden dimension.

  • Method Details

    • bindRequest

      public int bindRequest(int requestId)
      Binds a stable home row for requestId.
      Parameters:
      requestId - id aligned with DeltaNetStatePool.bindRequest(int).
      Returns:
      home row index.

    • unbindRequest

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

    • setRow

      public void setRow(int requestId, Tensor row)
      Writes row's content into requestId's durable anchor. Allocates the pool's backing tensor on the first call ever made on this pool (dtype/device taken from row); the anchor dimension itself is a fixed model hyperparameter, so no later reallocation is expected.
      Parameters:
      requestId - bound request id.
      row - anchor row [D] or [1, D].

    • hasRow

      public boolean hasRow(int requestId)
      Returns whether requestId has a written anchor row yet.
      Parameters:
      requestId - bound request id.
      Returns:
      true once setRow(int, Tensor) has been called for this id.

    • getRow

      public Tensor getRow(int requestId)
      Returns a fresh [D] copy of requestId's anchor row, or null if never written.
      Parameters:
      requestId - bound request id.
      Returns:
      owned [1, D] copy (matching the [B, D] shape QwenMtp.draftStep expects), or null.

    • getRows

      public Tensor getRows(int[] requestIds)
      Returns a fresh [B, D] tensor gathering each request's anchor row in order (batched MTP draft input), or null if never written.
      Parameters:
      requestIds - bound request ids (order = batch).
      Returns:
      owned [B, D] tensor, or null.

    • close

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