Record Class AttentionContext

java.lang.Object
java.lang.Record
smile.llm.attention.AttentionContext
Record Components:
scale - attention scale; <= 0 means kernel default (1/sqrt(D)).
dropout - dropout probability (torch_native only; usually 0 at inference).
isCausal - whether to apply causal masking when mask is null.
numQoHeads - query / output head count (GQA).
numKvHeads - key / value head count.
headDim - per-head dimension.
layerId - KV pool layer index (FlashInfer).
startPos - write position in the request (FlashInfer); uniform when startPositions is null.
seqLen - query sequence length this step.
cacheLen - total cached length after this step (startPos + seqLen) when cacheLens is null.
kvPool - shared pool, or null for contiguous torch_native-only calls.
kvMetadata - CSR page table for FlashInfer, or null.
workspace - FlashInfer workspace, or null.
startPositions - optional per-row write positions (decode); null = uniform.
cacheLens - optional per-row cache lengths (decode); null = uniform.
raggedIndptr - cumulative segment lengths for ragged contiguous prefill; null otherwise.
public record AttentionContext(double scale, double dropout, boolean isCausal, int numQoHeads, int numKvHeads, int headDim, int layerId, int startPos, int seqLen, int cacheLen, KvCachePool kvPool, FlashInferKvMetadata kvMetadata, FlashInferWorkspace workspace, int[] startPositions, int[] cacheLens, int[] raggedIndptr) extends Record
Per-call metadata for AttentionKernel.forward(Tensor, Tensor, Tensor, Tensor, AttentionContext).

  • Constructor Details

    • AttentionContext

      public AttentionContext(double scale, double dropout, boolean isCausal, int numQoHeads, int numKvHeads, int headDim, int layerId, int startPos, int seqLen, int cacheLen, KvCachePool kvPool, FlashInferKvMetadata kvMetadata, FlashInferWorkspace workspace, int[] startPositions, int[] cacheLens, int[] raggedIndptr)
      Creates an instance of a AttentionContext record class.
      Parameters:
      scale - the value for the scale record component
      dropout - the value for the dropout record component
      isCausal - the value for the isCausal record component
      numQoHeads - the value for the numQoHeads record component
      numKvHeads - the value for the numKvHeads record component
      headDim - the value for the headDim record component
      layerId - the value for the layerId record component
      startPos - the value for the startPos record component
      seqLen - the value for the seqLen record component
      cacheLen - the value for the cacheLen record component
      kvPool - the value for the kvPool record component
      kvMetadata - the value for the kvMetadata record component
      workspace - the value for the workspace record component
      startPositions - the value for the startPositions record component
      cacheLens - the value for the cacheLens record component
      raggedIndptr - the value for the raggedIndptr record component

  • Method Details

    • contiguous

      public static AttentionContext contiguous(double scale, double dropout, boolean isCausal)
      Contiguous SDPA context (gather path).
      Parameters:
      scale - attention scale; <= 0 means kernel default.
      dropout - dropout probability (usually 0 at inference).
      isCausal - whether to apply causal masking when mask is null.
      Returns:
      contiguous attention context.

    • ragged

      public static AttentionContext ragged(double scale, boolean isCausal, int numQoHeads, int numKvHeads, int headDim, int[] indptr)
      Ragged contiguous prefill (vision tower / varlen self-attention).
      Parameters:
      scale - attention scale; <= 0 means kernel default.
      isCausal - whether to apply causal masking when mask is null.
      numQoHeads - query / output head count (GQA).
      numKvHeads - key / value head count.
      headDim - per-head dimension.
      indptr - cumulative segment lengths [B+1] (CPU array; copied to device at call time).
      Returns:
      ragged contiguous attention context.

    • paged

      public static AttentionContext paged(double scale, boolean isCausal, int numQoHeads, int numKvHeads, int headDim, int layerId, int startPos, int seqLen, int cacheLen, KvCachePool kvPool, FlashInferKvMetadata kvMetadata, FlashInferWorkspace workspace)
      FlashInfer / paged context (uniform length).
      Parameters:
      scale - attention scale; <= 0 means kernel default.
      isCausal - whether to apply causal masking when mask is null.
      numQoHeads - query / output head count (GQA).
      numKvHeads - key / value head count.
      headDim - per-head dimension.
      layerId - KV pool layer index.
      startPos - write position in the request.
      seqLen - query sequence length this step.
      cacheLen - total cached length after this step.
      kvPool - shared KV cache pool.
      kvMetadata - CSR page table for FlashInfer.
      workspace - FlashInfer workspace.
      Returns:
      paged attention context.

    • pagedRagged

      public static AttentionContext pagedRagged(double scale, boolean isCausal, int numQoHeads, int numKvHeads, int headDim, int layerId, int seqLen, int[] startPositions, int[] cacheLens, KvCachePool kvPool, FlashInferKvMetadata kvMetadata, FlashInferWorkspace workspace)
      FlashInfer / paged context with per-row decode lengths / positions.
      Parameters:
      scale - attention scale; <= 0 means kernel default.
      isCausal - whether to apply causal masking when mask is null.
      numQoHeads - query / output head count (GQA).
      numKvHeads - key / value head count.
      headDim - per-head dimension.
      layerId - KV pool layer index.
      seqLen - query sequence length this step.
      startPositions - optional per-row write positions; null = uniform.
      cacheLens - optional per-row cache lengths; null = uniform.
      kvPool - shared KV cache pool.
      kvMetadata - CSR page table for FlashInfer.
      workspace - FlashInfer workspace.
      Returns:
      paged ragged attention context.

    • isPaged

      public boolean isPaged()
      Returns whether this call carries paged-KV metadata.
      Returns:
      true when this call carries paged-KV metadata.

    • isRagged

      public boolean isRagged()
      Returns whether per-row cache lengths are set.
      Returns:
      true when per-row cache lengths are set.

    • isRaggedContiguous

      public boolean isRaggedContiguous()
      Returns whether this is FlashInfer ragged contiguous Q/K/V (non-paged).
      Returns:
      true for FlashInfer ragged contiguous Q/K/V (non-paged).

    • toString

      public final String toString()
      Returns a string representation of this record class. The representation contains the name of the class, followed by the name and value of each of the record components.
      Specified by:
      toString in class Record
      Returns:
      a string representation of this object

    • hashCode

      public final int hashCode()
      Returns a hash code value for this object. The value is derived from the hash code of each of the record components.
      Specified by:
      hashCode in class Record
      Returns:
      a hash code value for this object

    • equals

      public final boolean equals(Object o)
      Indicates whether some other object is "equal to" this one. The objects are equal if the other object is of the same class and if all the record components are equal. Reference components are compared with Objects::equals(Object,Object); primitive components are compared with the compare method from their corresponding wrapper classes.
      Specified by:
      equals in class Record
      Parameters:
      o - the object with which to compare
      Returns:
      true if this object is the same as the o argument; false otherwise.

    • scale

      public double scale()
      Returns the value of the scale record component.
      Returns:
      the value of the scale record component

    • dropout

      public double dropout()
      Returns the value of the dropout record component.
      Returns:
      the value of the dropout record component

    • isCausal

      public boolean isCausal()
      Returns the value of the isCausal record component.
      Returns:
      the value of the isCausal record component

    • numQoHeads

      public int numQoHeads()
      Returns the value of the numQoHeads record component.
      Returns:
      the value of the numQoHeads record component

    • numKvHeads

      public int numKvHeads()
      Returns the value of the numKvHeads record component.
      Returns:
      the value of the numKvHeads record component

    • headDim

      public int headDim()
      Returns the value of the headDim record component.
      Returns:
      the value of the headDim record component

    • layerId

      public int layerId()
      Returns the value of the layerId record component.
      Returns:
      the value of the layerId record component

    • startPos

      public int startPos()
      Returns the value of the startPos record component.
      Returns:
      the value of the startPos record component

    • seqLen

      public int seqLen()
      Returns the value of the seqLen record component.
      Returns:
      the value of the seqLen record component

    • cacheLen

      public int cacheLen()
      Returns the value of the cacheLen record component.
      Returns:
      the value of the cacheLen record component

    • kvPool

      public KvCachePool kvPool()
      Returns the value of the kvPool record component.
      Returns:
      the value of the kvPool record component

    • kvMetadata

      public FlashInferKvMetadata kvMetadata()
      Returns the value of the kvMetadata record component.
      Returns:
      the value of the kvMetadata record component

    • workspace

      public FlashInferWorkspace workspace()
      Returns the value of the workspace record component.
      Returns:
      the value of the workspace record component

    • startPositions

      public int[] startPositions()
      Returns the value of the startPositions record component.
      Returns:
      the value of the startPositions record component

    • cacheLens

      public int[] cacheLens()
      Returns the value of the cacheLens record component.
      Returns:
      the value of the cacheLens record component

    • raggedIndptr

      public int[] raggedIndptr()
      Returns the value of the raggedIndptr record component.
      Returns:
      the value of the raggedIndptr record component