attention

package
v1.58.0 Latest Latest
Warning

This package is not in the latest version of its module.

 Go to latest Published: Aug 8, 2026 License: Apache-2.0 Imports: 18 Imported by: 1

Documentation

Overview

Package attention provides attention mechanisms for neural networks.

Stability: stable

Package attention exposes a graph.Node wrapper around ScaledDotProductAttention so consumers (e.g., Wolf cross-attention) can compose fused SDPA via graph.Builder without duplicating the underlying math.

Package attention provides attention mechanisms for neural networks.

Package attention provides attention mechanisms for neural networks.

Index

Constants

This section is empty.

Variables

This section is empty.

Functions

func BuildCausalSlidingWindowMask added in v0.2.1 

func BuildCausalSlidingWindowMask[T tensor.Numeric](seqLen, windowSize int) *tensor.TensorNumeric[T]

BuildCausalSlidingWindowMask creates a causal attention mask that also restricts attention to the last windowSize positions. Positions outside the window or in the future are set to a large negative value. Shape: [1, 1, seqLen, seqLen].

func BuildGlobalAttention added in v0.2.1 

func BuildGlobalAttention[T tensor.Numeric](
	engine compute.Engine[T],
	ops numeric.Arithmetic[T],
	_ string,
	_ map[string]*graph.Parameter[T],
	attributes map[string]interface{},
) (graph.Node[T], error)

BuildGlobalAttention constructs a GlobalAttention node from attributes. Required attributes: - "embed_dim" (int): embedding dimension - "num_heads" (int): number of attention heads - "num_kv_heads" (int): number of key-value heads

func BuildGroupQueryAttention added in v0.2.0 

func BuildGroupQueryAttention[T tensor.Numeric](
	engine compute.Engine[T],
	ops numeric.Arithmetic[T],
	name string,
	params map[string]*graph.Parameter[T],
	attributes map[string]interface{},
) (graph.Node[T], error)

BuildGroupQueryAttention constructs a GroupedQueryAttention node for the model builder. Unused parameters are accepted to satisfy the common builder signature.

func BuildMultiHeadLatentAttention added in v0.2.1 

func BuildMultiHeadLatentAttention[T tensor.Numeric](
	engine compute.Engine[T],
	ops numeric.Arithmetic[T],
	name string,
	params map[string]*graph.Parameter[T],
	attributes map[string]any,
) (graph.Node[T], error)

BuildMultiHeadLatentAttention constructs a MultiHeadLatentAttention node for the model builder. It reads kv_lora_dim, num_heads, head_dim, and max_seq_len from attributes, and loads W_Q, W_DKV, W_UK, W_UV, W_O from node parameters.

func BuildNativeSparseAttention added in v1.29.0 

func BuildNativeSparseAttention[T tensor.Numeric](
	engine compute.Engine[T],
	ops numeric.Arithmetic[T],
	name string,
	params map[string]*graph.Parameter[T],
	attributes map[string]interface{},
) (graph.Node[T], error)

BuildNativeSparseAttention constructs a NativeSparseAttention node for the model builder. It reads model_dim, num_heads, num_kv_heads, block_size, top_blocks, top_tokens, and window_size from attributes, and loads gate_coarse, gate_fine, gate_window from node parameters.

func BuildSparseRoutedAttention added in v1.29.0 

func BuildSparseRoutedAttention[T tensor.Numeric](
	engine compute.Engine[T],
	ops numeric.Arithmetic[T],
	name string,
	params map[string]*graph.Parameter[T],
	attributes map[string]interface{},
) (graph.Node[T], error)

BuildSparseRoutedAttention constructs a SparseRoutedAttention node for the model builder. It reads num_heads, num_kv_heads, head_dim, segment_size, top_k, max_seq_len, and rope_base from attributes. A nil KV cache is used since cache binding happens at generation time.

func QKNorm added in v0.2.0 

func QKNorm[T tensor.Numeric](ctx context.Context, engine compute.Engine[T], q, k *tensor.TensorNumeric[T], epsilon float64) (qNorm, kNorm *tensor.TensorNumeric[T], err error)

QKNorm applies a form of normalization to Query (Q) and Key (K) tensors to stabilize attention score scales, similar to RMSNorm. It normalizes Q and K independently by their respective RMS values. All operations use Engine primitives so they appear in the ExecutionPlan instruction tape.

Types

type AttentionHead added in v0.2.0 

type AttentionHead[T tensor.Numeric] struct {
	// contains filtered or unexported fields
}

AttentionHead implements a single attention head, including linear projections for Query, Key, and Value, followed by scaled dot-product attention.

func NewAttentionHead added in v0.2.0 

func NewAttentionHead[T tensor.Numeric](engine compute.Engine[T], inputDim, headDim int, opts ...AttentionHeadOption[T]) (*AttentionHead[T], error)

NewAttentionHead creates a new AttentionHead instance. inputDim is the dimension of the input features. headDim is the dimension of the query, key, and value vectors for this head.

func (*AttentionHead[T]) Attributes added in v0.2.1 

func (ah *AttentionHead[T]) Attributes() map[string]interface{}

Attributes returns the attributes for the AttentionHead.

func (*AttentionHead[T]) Backward added in v0.2.0 

func (ah *AttentionHead[T]) Backward(ctx context.Context, mode types.BackwardMode, dOut *tensor.TensorNumeric[T], inputs ...*tensor.TensorNumeric[T]) ([]*tensor.TensorNumeric[T], error)

Backward computes the gradients for the AttentionHead. dOut has shape (batch, seq_len, head_dim). inputs[0] has shape (batch, seq_len, input_dim).

func (*AttentionHead[T]) Forward added in v0.2.0 

func (ah *AttentionHead[T]) Forward(ctx context.Context, inputs ...*tensor.TensorNumeric[T]) (*tensor.TensorNumeric[T], error)

Forward computes the output of the attention head. input is expected to be a 3D tensor (batch_size, seq_len, input_dim).

func (*AttentionHead[T]) OpType added in v0.2.1 

func (ah *AttentionHead[T]) OpType() string

OpType returns the operation type of the AttentionHead.

func (*AttentionHead[T]) OutputShape added in v0.2.0 

func (ah *AttentionHead[T]) OutputShape() []int

OutputShape returns the output shape of the AttentionHead. It assumes the input shape is (batch_size, seq_len, input_dim). The output shape will be (batch_size, seq_len, head_dim).

func (*AttentionHead[T]) Parameters added in v0.2.0 

func (ah *AttentionHead[T]) Parameters() []*graph.Parameter[T]

Parameters returns all trainable parameters of the AttentionHead.

func (*AttentionHead[T]) SetSaver added in v1.49.0 

func (ah *AttentionHead[T]) SetSaver(sv graph.Saver[T])

SetSaver implements graph.SaverAware, fanning the graph-provided Saver into the composed SDPA so its cached q/k/v/attention-weights tensors are pinned until this node's Backward returns (ztensor ADR 006).

type AttentionHeadOption added in v0.2.0 

type AttentionHeadOption[T tensor.Numeric] func(*AttentionHeadOptions[T])

AttentionHeadOption applies an option to AttentionHeadOptions.

func WithBidirectionalAttention added in v1.9.0 

func WithBidirectionalAttention[T tensor.Numeric]() AttentionHeadOption[T]

WithBidirectionalAttention returns an option that disables causal masking, allowing every token to attend to every other token. This is used by encoder-style models such as BERT.

type AttentionHeadOptions added in v0.2.0 

type AttentionHeadOptions[T tensor.Numeric] struct {
	// contains filtered or unexported fields
}

AttentionHeadOptions holds configuration options for AttentionHead.

type BlockTableReader added in v0.2.1 

type BlockTableReader[T tensor.Numeric] interface {
	// ReadKV returns the cached key and value tensors for the given layer
	// as contiguous [batch*numKVHeads, seqLen, headDim] tensors read directly
	// from blocks. Returns false if the layer has no cached data.
	ReadKV(layer int) (k, v *tensor.TensorNumeric[T], ok bool)
}

BlockTableReader reads key/value tensors directly from paged block tables, avoiding the gather-to-contiguous copy. Implementations should iterate over blocks and return the full KV sequence for a given layer.

type FusedSDPA added in v1.49.0 

type FusedSDPA[T tensor.Numeric] struct {
	// contains filtered or unexported fields
}

FusedSDPA wraps ScaledDotProductAttention as a graph.Node[T]. It accepts (Q, K, V) or (Q, K, V, mask) inputs in Forward and delegates to the inner SDPA. Backward delegates to the inner SDPA and pads a nil gradient slot for the mask input when one was supplied so input/grad indexing stays aligned.

func NewFusedSDPA added in v1.49.0 

func NewFusedSDPA[T tensor.Numeric](engine compute.Engine[T], headDim int, opts ...FusedSDPAOption[T]) *FusedSDPA[T]

NewFusedSDPA constructs a FusedSDPA graph.Node wrapping a ScaledDotProductAttention configured with the given engine, head dimension, and options.

func NewFusedSDPAFrom added in v1.49.0 

func NewFusedSDPAFrom[T tensor.Numeric](sdpa *ScaledDotProductAttention[T]) *FusedSDPA[T]

NewFusedSDPAFrom wraps an existing ScaledDotProductAttention as a graph.Node. Useful when callers want to share state with an already-configured layer.

func (*FusedSDPA[T]) Attributes added in v1.49.0 

func (n *FusedSDPA[T]) Attributes() map[string]interface{}

Attributes implements graph.Node and reports configuration knobs needed to reconstruct the op (head dimension and causal flag).

func (*FusedSDPA[T]) Backward added in v1.49.0 

func (n *FusedSDPA[T]) Backward(ctx context.Context, mode types.BackwardMode, outputGradient *tensor.TensorNumeric[T], inputs ...*tensor.TensorNumeric[T]) ([]*tensor.TensorNumeric[T], error)

Backward delegates to the inner SDPA. The inner returns gradients for [Q, K, V]; we append a nil gradient for mask when one was supplied so the returned slice aligns 1:1 with the inputs passed to Forward.

func (*FusedSDPA[T]) Forward added in v1.49.0 

func (n *FusedSDPA[T]) Forward(ctx context.Context, inputs ...*tensor.TensorNumeric[T]) (*tensor.TensorNumeric[T], error)

Forward delegates to the inner SDPA. inputs must be (Q, K, V) or (Q, K, V, mask).

func (*FusedSDPA[T]) OpType added in v1.49.0 

func (n *FusedSDPA[T]) OpType() string

OpType implements graph.Node.

func (*FusedSDPA[T]) OutputShape added in v1.49.0 

func (n *FusedSDPA[T]) OutputShape() []int

OutputShape returns the cached shape of the most recent Forward, or nil.

func (*FusedSDPA[T]) Parameters added in v1.49.0 

func (n *FusedSDPA[T]) Parameters() []*graph.Parameter[T]

Parameters implements graph.Node. SDPA has no trainable parameters of its own.

func (*FusedSDPA[T]) SetSaver added in v1.49.1 

func (n *FusedSDPA[T]) SetSaver(s graph.Saver[T])

SetSaver implements graph.SaverAware by fanning the graph's Saver into the inner SDPA. The inner SDPA caches Q/K/V and the attention weights in Forward and consumes them in Backward (which is called with nil inputs), so they are SAVE-class intermediates: without this wiring they would be unpinned arena intermediates on GPU engines, and downstream forward ops could overwrite them before Backward runs (zerfoo#864, the #842 class).

type FusedSDPAOption added in v1.49.0 

type FusedSDPAOption[T tensor.Numeric] func(*fusedSDPAOptions[T])

FusedSDPAOption configures a FusedSDPA at construction time.

func WithFusedSDPABidirectional added in v1.49.0 

func WithFusedSDPABidirectional[T tensor.Numeric]() FusedSDPAOption[T]

WithFusedSDPABidirectional disables causal masking (encoder-style attention).

func WithFusedSDPAHeadCounts added in v1.49.0 

func WithFusedSDPAHeadCounts[T tensor.Numeric](numQueryHeads, numKVHeads int) FusedSDPAOption[T]

WithFusedSDPAHeadCounts sets query/KV head counts so the inner SDPA can dispatch to the split-KV flash decode kernel where applicable.

type GQAOption added in v0.2.0 

type GQAOption[T tensor.Numeric] func(*GQAOptions[T])

GQAOption is a function that applies an option to GQAOptions.

func WithBidirectionalGQA added in v1.9.0 

func WithBidirectionalGQA[T tensor.Numeric]() GQAOption[T]

WithBidirectionalGQA returns an option that disables causal masking in the grouped query attention layer, allowing every position to attend to every other position. This is required for encoder-style models such as BERT.

func WithExternalKV added in v1.49.0 

func WithExternalKV[T tensor.Numeric]() GQAOption[T]

WithExternalKV returns an option that configures the GroupedQueryAttention layer to read K and V from additional Forward inputs instead of computing them from internal wk/wv projections. When set, the constructor does not instantiate wk, wv, or k_norm; Forward expects inputs[1] to be K and inputs[2] to be V, each with shape [batch, seq, numKeyValueHeads * headDim] (matching what the internal projection would produce). RoPE for K still applies (if configured). Q, q_norm, w_out, attention math, and output projection remain unchanged. This option is mutually exclusive with the kEqV mode configured via SetKEqV. See ADR-087 for rationale.

func WithMaxSeqLen added in v0.2.0 

func WithMaxSeqLen[T tensor.Numeric](maxSeqLen int) GQAOption[T]

WithMaxSeqLen sets the maximum sequence length for Rotary Positional Embeddings.

func WithNoRoPE added in v1.32.0 

func WithNoRoPE[T tensor.Numeric]() GQAOption[T]

WithNoRoPE returns an option that disables rotary positional embeddings. Models like GPT-2 use learned position embeddings instead of RoPE, so the GQA layer should pass Q and K through without rotational encoding.

func WithRopeBase added in v0.2.0 

func WithRopeBase[T tensor.Numeric](base float64) GQAOption[T]

WithRopeBase sets the base for Rotary Positional Embeddings.

type GQAOptions added in v0.2.0 

type GQAOptions[T tensor.Numeric] struct {
	Base          float64
	MaxSeqLen     int
	Bidirectional bool // when true, disables causal masking for encoder-style models
	NoRoPE        bool // when true, skip RoPE creation (for models like GPT-2 that use learned position embeddings)
	ExternalKV    bool // when true, K/V are supplied as Forward inputs; wk/wv/k_norm are not instantiated
}

GQAOptions holds configuration options for the GroupedQueryAttention layer.

type GlobalAttention added in v0.2.0 

type GlobalAttention[T tensor.Numeric] struct {
	// contains filtered or unexported fields
}

GlobalAttention wraps GroupedQueryAttention to provide a global attention interface.

func NewGlobalAttention added in v0.2.0 

func NewGlobalAttention[T tensor.Numeric](
	engine compute.Engine[T],
	ops numeric.Arithmetic[T],
	modelDim, numQueryHeads, numKeyValueHeads int,
	options ...GlobalAttentionOption,
) (*GlobalAttention[T], error)

NewGlobalAttention creates a new GlobalAttention layer.

Parameters: - engine: compute engine for tensor operations - ops: arithmetic operations for the numeric type - modelDim: model dimension - numQueryHeads: number of query heads - numKeyValueHeads: number of key/value heads - options: functional options for configuration

Default values: - base: 10000.0 - maxSeqLen: 2048.

func NewGlobalAttentionFromParams added in v0.2.0 

func NewGlobalAttentionFromParams[T tensor.Numeric](gqa *GroupedQueryAttention[T]) *GlobalAttention[T]

NewGlobalAttentionFromParams creates a new GlobalAttention layer from an existing GroupedQueryAttention layer.

func (*GlobalAttention[T]) Attributes added in v0.2.0 

func (ga *GlobalAttention[T]) Attributes() map[string]interface{}

Attributes returns the attributes.

func (*GlobalAttention[T]) Backward added in v0.2.0 

func (ga *GlobalAttention[T]) Backward(ctx context.Context, mode types.BackwardMode, dOut *tensor.TensorNumeric[T], inputs ...*tensor.TensorNumeric[T]) ([]*tensor.TensorNumeric[T], error)

Backward delegates the backward pass to the wrapped GroupedQueryAttention.

func (*GlobalAttention[T]) Forward added in v0.2.0 

func (ga *GlobalAttention[T]) Forward(ctx context.Context, inputs ...*tensor.TensorNumeric[T]) (*tensor.TensorNumeric[T], error)

Forward computes the forward pass of the GlobalAttention layer.

func (*GlobalAttention[T]) OpType added in v0.2.0 

func (ga *GlobalAttention[T]) OpType() string

OpType returns the operation type.

func (*GlobalAttention[T]) OutputShape added in v0.2.0 

func (ga *GlobalAttention[T]) OutputShape() []int

OutputShape returns the output shape of the GlobalAttention layer.

func (*GlobalAttention[T]) Parameters added in v0.2.0 

func (ga *GlobalAttention[T]) Parameters() []*graph.Parameter[T]

Parameters returns the parameters of the GlobalAttention layer.

func (*GlobalAttention[T]) ScaleRope added in v0.2.0 

func (ga *GlobalAttention[T]) ScaleRope(ctx context.Context, factor float64) error

ScaleRope scales the rotary positional embeddings.

func (*GlobalAttention[T]) SetLayerIndex added in v0.2.1 

func (ga *GlobalAttention[T]) SetLayerIndex(idx int)

SetLayerIndex sets the layer index for KV cache routing.

func (*GlobalAttention[T]) SetSaver added in v1.49.0 

func (ga *GlobalAttention[T]) SetSaver(sv graph.Saver[T])

SetSaver implements graph.SaverAware, fanning the Saver into the wrapped GroupedQueryAttention (ztensor ADR 006).

type GlobalAttentionOption added in v0.2.0 

type GlobalAttentionOption func(*GlobalAttentionOptions)

GlobalAttentionOption is a function that configures GlobalAttentionOptions.

func WithGlobalAttentionBase added in v0.2.0 

func WithGlobalAttentionBase(base float64) GlobalAttentionOption

WithGlobalAttentionBase sets the base (theta) parameter for rotary positional embeddings.

func WithGlobalAttentionMaxSeqLen added in v0.2.0 

func WithGlobalAttentionMaxSeqLen(maxSeqLen int) GlobalAttentionOption

WithGlobalAttentionMaxSeqLen sets the maximum sequence length.

type GlobalAttentionOptions added in v0.2.0 

type GlobalAttentionOptions struct {
	Base      float64
	MaxSeqLen int
}

GlobalAttentionOptions holds configuration options for GlobalAttention layer.

type GroupedQueryAttention 

type GroupedQueryAttention[T tensor.Numeric] struct {

	// LayerIndex identifies this layer within a model for KV cache indexing.
	LayerIndex int

	// SlidingWindowSize, if > 0, restricts attention to the last N positions
	// using a causal sliding window mask during prefill (seqLen > 1).
	SlidingWindowSize int
	// contains filtered or unexported fields
}

GroupedQueryAttention implements grouped query attention mechanism.

func NewGroupedQueryAttention 

func NewGroupedQueryAttention[T tensor.Numeric](
	engine compute.Engine[T],
	ops numeric.Arithmetic[T],
	modelDim, numQueryHeads, numKeyValueHeads int,
	opts ...GQAOption[T],
) (*GroupedQueryAttention[T], error)

NewGroupedQueryAttention creates a new GroupedQueryAttention layer. modelDim: The dimension of the input and output of the block (d_model). numQueryHeads: The number of query heads. numKeyValueHeads: The number of key/value heads.

func NewGroupedQueryAttentionFromParams added in v0.2.0 

func NewGroupedQueryAttentionFromParams[T tensor.Numeric](
	engine compute.Engine[T],
	ops numeric.Arithmetic[T],
	modelDim, numQueryHeads, numKeyValueHeads int,
	wq, wk, wv, wo *core.Dense[T],
	rope *embeddings.RotaryPositionalEmbedding[T],
	headDimOverride ...int,
) (*GroupedQueryAttention[T], error)

NewGroupedQueryAttentionFromParams creates a new GroupedQueryAttention layer from existing parameters. headDimOverride, if > 0, sets the per-head dimension explicitly instead of deriving it from modelDim/numQueryHeads. This is required for architectures like Gemma 3 where key_length differs from hidden_size/num_heads.

func (*GroupedQueryAttention[T]) Attributes added in v0.2.0 

func (gqa *GroupedQueryAttention[T]) Attributes() map[string]interface{}

Attributes returns the attributes.

func (*GroupedQueryAttention[T]) Backward 

func (gqa *GroupedQueryAttention[T]) Backward(ctx context.Context, mode types.BackwardMode, dOut *tensor.TensorNumeric[T], inputs ...*tensor.TensorNumeric[T]) ([]*tensor.TensorNumeric[T], error)

Backward computes the gradients for GroupedQueryAttention.

The backward mirrors the forward in reverse order:

  1. wo backward
  2. Reverse reshape/transpose (head concatenation)
  3. SDPA backward
  4. Reverse K/V head replication (sum over group copies)
  5. RoPE backward
  6. Reverse head split (reshape/transpose back to projection shape)
  7. wq/wk/wv backward

func (*GroupedQueryAttention[T]) ExternalKV added in v1.49.0 

func (gqa *GroupedQueryAttention[T]) ExternalKV() bool

ExternalKV reports whether this layer was configured with WithExternalKV.

func (*GroupedQueryAttention[T]) Forward 

func (gqa *GroupedQueryAttention[T]) Forward(ctx context.Context, inputs ...*tensor.TensorNumeric[T]) (*tensor.TensorNumeric[T], error)

Forward computes the grouped query attention.

func (*GroupedQueryAttention[T]) KPort added in v1.49.0 

func (gqa *GroupedQueryAttention[T]) KPort() graph.Node[T]

KPort returns a graph.Node[T] that, when executed, yields the layer's most-recently-computed K in shape [batch, numKVHeads, seqLen, headDim]. In external-K/V mode, the captured K is whatever tensor was supplied via Forward's inputs[1]; otherwise it is the layer's own projection output after optional per-head K norm and post-projection RoPE.

Donor -> shared wiring is performed by the caller: at graph-build time the builder wires donor.KPort() as an input producing K, then passes the resulting tensor into the shared layer's Forward via inputs[1]. See ADR-087.

The returned node's Forward must be called AFTER the owner GQA's Forward pass -- the port re-exports the cached tensor produced there. Calling Forward on the port before the owner has produced K yields a descriptive error (see gqaKVPort.Forward). The node carries no parameters and has no independent gradient.

func (*GroupedQueryAttention[T]) MergedQKVParameter added in v0.2.1 

func (gqa *GroupedQueryAttention[T]) MergedQKVParameter() *graph.Parameter[T]

MergedQKVParameter returns the merged QKV parameter for GPU upload, or nil if not set.

func (*GroupedQueryAttention[T]) OpType added in v0.2.0 

func (gqa *GroupedQueryAttention[T]) OpType() string

OpType returns the operation type.

func (*GroupedQueryAttention[T]) OutputShape 

func (gqa *GroupedQueryAttention[T]) OutputShape() []int

OutputShape returns the output shape of the GroupedQueryAttention.

func (*GroupedQueryAttention[T]) Parameters 

func (gqa *GroupedQueryAttention[T]) Parameters() []*graph.Parameter[T]

Parameters returns the parameters of the GroupedQueryAttention layer.

func (*GroupedQueryAttention[T]) ScaleRope added in v0.2.0 

func (gqa *GroupedQueryAttention[T]) ScaleRope(ctx context.Context, factor float64) error

ScaleRope scales the rotary positional embeddings.

func (*GroupedQueryAttention[T]) SetBidirectional added in v1.9.0 

func (gqa *GroupedQueryAttention[T]) SetBidirectional(bidirectional bool)

SetBidirectional enables or disables bidirectional (non-causal) attention.

func (*GroupedQueryAttention[T]) SetBlockTableReader added in v0.2.1 

func (gqa *GroupedQueryAttention[T]) SetBlockTableReader(r BlockTableReader[T])

SetBlockTableReader sets an optional BlockTableReader that provides KV data directly from paged block tables, bypassing the standard cache gather path.

func (*GroupedQueryAttention[T]) SetDocumentBoundaries added in v1.28.0 

func (gqa *GroupedQueryAttention[T]) SetDocumentBoundaries(boundaries []int)

SetDocumentBoundaries sets document boundary positions for document-wise RoPE. When boundaries are set, position IDs reset to 0 at each boundary so each document receives independent positional encoding during multi-document inference. Boundaries are sequence positions (0-indexed) where new documents begin. Pass nil to disable document-wise mode.

func (*GroupedQueryAttention[T]) SetExternalKV added in v1.49.0 

func (gqa *GroupedQueryAttention[T]) SetExternalKV(v bool)

SetExternalKV configures the layer to read K/V from Forward inputs[1]/[2] instead of projecting from its own wk/wv weights. When enabled, wk, wv, and kNorm are cleared (the shared-KV layer does not use them per ADR-087 and HuggingFace transformers modeling_gemma4.py line 1167 "Layers sharing kv states don't need any weight matrices"). Mutually exclusive with kEqV.

This setter is the FromParams-path counterpart to the WithExternalKV constructor option: builders that load weights from GGUF construct GQA via NewGroupedQueryAttentionFromParams with wk=nil, wv=nil (or disposable wk/wv), then invoke SetExternalKV(true) to flip the forward path. Idempotent; passing false re-enables internal projection iff wk/wv are still populated.

func (*GroupedQueryAttention[T]) SetKEqV added in v1.49.0 

func (gqa *GroupedQueryAttention[T]) SetKEqV(v bool)

SetKEqV configures GQA to use the K projection output for both K and V. When enabled, the V projection (wv) is skipped and K output is used as V. This implements Gemma 4's unified K=V projection for global attention layers. Mutually exclusive with WithExternalKV: panics if both are requested.

func (*GroupedQueryAttention[T]) SetMergedQKV added in v0.2.1 

func (gqa *GroupedQueryAttention[T]) SetMergedQKV(weight *tensor.TensorNumeric[T], qDim, kDim, vDim int)

SetMergedQKV sets a merged Q/K/V weight tensor for single-GEMV decode optimization. During decode (seqLen=1), a single MatMul with this weight replaces three separate Q/K/V projections, reducing kernel launch overhead. The output is split into Q, K, V using zero-copy GPU storage views.

func (*GroupedQueryAttention[T]) SetQKNormPreReshape added in v1.36.0 

func (gqa *GroupedQueryAttention[T]) SetQKNormPreReshape(v bool)

SetQKNormPreReshape configures whether QK norms are applied before the head reshape. Set true for architectures (e.g. MiniMax-M2) whose norm weights are [nH*hD] rather than [hD].

func (*GroupedQueryAttention[T]) SetQKNormWeights added in v0.2.1 

func (gqa *GroupedQueryAttention[T]) SetQKNormWeights(qWeight, kWeight *tensor.TensorNumeric[T], eps float32)

SetQKNormWeights stores raw RMSNorm weights for the fused QK norm+RoPE decode path. When set alongside SetQKNorms, the fused kernel replaces 4 kernel launches (Q norm, K norm, Q RoPE, K RoPE) with 1 during decode.

func (*GroupedQueryAttention[T]) SetQKNorms added in v0.2.1 

func (gqa *GroupedQueryAttention[T]) SetQKNorms(qNorm, kNorm graph.Node[T])

SetQKNorms sets optional per-head RMSNorm layers for Q and K projections. Used by architectures like Gemma 3 that normalize Q/K after projection. When externalKV is set, the kNorm argument is ignored (K is supplied pre-computed and any normalization must have been applied by the donor).

func (*GroupedQueryAttention[T]) SetSaver added in v1.49.0 

func (gqa *GroupedQueryAttention[T]) SetSaver(sv graph.Saver[T])

SetSaver implements graph.SaverAware, fanning the Saver into the composed SDPA so its cached q/k/v/attention-weights tensors are pinned until this node's Backward returns (ztensor ADR 006).

func (*GroupedQueryAttention[T]) VPort added in v1.49.0 

func (gqa *GroupedQueryAttention[T]) VPort() graph.Node[T]

VPort returns a graph.Node[T] that yields the layer's most-recently- computed V in shape [batch, numKVHeads, seqLen, headDim]. See KPort for semantics, donor-wiring conventions, and nil-safety contract.

type LocalAttention added in v0.2.0 

type LocalAttention[T tensor.Numeric] struct {
	// contains filtered or unexported fields
}

LocalAttention implements a local, sliding-window self-attention mechanism.

func NewLocalAttention added in v0.2.0 

func NewLocalAttention[T tensor.Numeric](
	engine compute.Engine[T],
	ops numeric.Arithmetic[T],
	modelDim, numQueryHeads, numKeyValueHeads, windowSize int,
	opts ...LocalAttentionOption[T],
) (*LocalAttention[T], error)

NewLocalAttention creates a new LocalAttention layer.

func (*LocalAttention[T]) Attributes added in v0.2.1 

func (la *LocalAttention[T]) Attributes() map[string]interface{}

Attributes returns the attributes of the LocalAttention layer.

func (*LocalAttention[T]) Backward added in v0.2.0 

func (la *LocalAttention[T]) Backward(ctx context.Context, mode types.BackwardMode, dOut *tensor.TensorNumeric[T], inputs ...*tensor.TensorNumeric[T]) ([]*tensor.TensorNumeric[T], error)

Backward delegates the backward pass to the wrapped GroupedQueryAttention.

func (*LocalAttention[T]) Forward added in v0.2.0 

func (la *LocalAttention[T]) Forward(ctx context.Context, inputs ...*tensor.TensorNumeric[T]) (*tensor.TensorNumeric[T], error)

Forward computes the forward pass of the LocalAttention layer.

func (*LocalAttention[T]) OpType added in v0.2.1 

func (la *LocalAttention[T]) OpType() string

OpType returns the operation type of the LocalAttention layer.

func (*LocalAttention[T]) OutputShape added in v0.2.0 

func (la *LocalAttention[T]) OutputShape() []int

OutputShape returns the output shape of the LocalAttention layer.

func (*LocalAttention[T]) Parameters added in v0.2.0 

func (la *LocalAttention[T]) Parameters() []*graph.Parameter[T]

Parameters returns the parameters of the LocalAttention layer.

func (*LocalAttention[T]) SetSaver added in v1.49.0 

func (la *LocalAttention[T]) SetSaver(sv graph.Saver[T])

SetSaver implements graph.SaverAware, fanning the Saver into the wrapped GroupedQueryAttention (ztensor ADR 006).

type LocalAttentionOption added in v0.2.0 

type LocalAttentionOption[T tensor.Numeric] func(*LocalAttentionOptions[T])

LocalAttentionOption is a function that applies an option to LocalAttentionOptions.

func WithLocalMaxSeqLen added in v0.2.0 

func WithLocalMaxSeqLen[T tensor.Numeric](maxSeqLen int) LocalAttentionOption[T]

WithLocalMaxSeqLen sets the maximum sequence length for Rotary Positional Embeddings.

maxSeqLen: The maximum sequence length for Rotary Positional Embeddings.

func WithLocalRopeBase added in v0.2.0 

func WithLocalRopeBase[T tensor.Numeric](base float64) LocalAttentionOption[T]

WithLocalRopeBase sets the base for Rotary Positional Embeddings.

base: The base for Rotary Positional Embeddings.

type LocalAttentionOptions added in v0.2.0 

type LocalAttentionOptions[T tensor.Numeric] struct {
	Base      float64
	MaxSeqLen int
}

LocalAttentionOptions holds configuration options for the LocalAttention layer.

type MultiHeadLatentAttention added in v0.2.1 

type MultiHeadLatentAttention[T tensor.Numeric] struct {
	// contains filtered or unexported fields
}

MultiHeadLatentAttention implements Multi-head Latent Attention (MLA) as used in DeepSeek V3/R1. MLA compresses KV into a low-rank latent vector, dramatically reducing KV cache size.

Partial RoPE: RoPE is applied only to the first ropeHeadDim dimensions of Q and K. The remaining dimensions are position-independent, matching the DeepSeek V3 paper specification.

func NewMultiHeadLatentAttention added in v0.2.1 

func NewMultiHeadLatentAttention[T tensor.Numeric](
	engine compute.Engine[T],
	ops numeric.Arithmetic[T],
	numHeads, headDim, kvLoraDim, ropeHeadDim int,
	wQ, wDKV, wUK, wUV, wO *core.Dense[T],
	rope *embeddings.RotaryPositionalEmbedding[T],
) *MultiHeadLatentAttention[T]

NewMultiHeadLatentAttention creates a new MLA layer. ropeHeadDim specifies how many of the headDim dimensions receive RoPE. If ropeHeadDim <= 0 or >= headDim, RoPE is applied to all dimensions (backwards-compatible behavior).

func (*MultiHeadLatentAttention[T]) Attributes added in v0.2.1 

func (m *MultiHeadLatentAttention[T]) Attributes() map[string]any

Attributes returns the layer attributes.

func (*MultiHeadLatentAttention[T]) Backward added in v0.2.1 

func (m *MultiHeadLatentAttention[T]) Backward(ctx context.Context, mode types.BackwardMode, dOut *tensor.TensorNumeric[T], inputs ...*tensor.TensorNumeric[T]) ([]*tensor.TensorNumeric[T], error)

Backward computes gradients for MLA. dOut: [batch, seqLen, hidden], inputs: original input to Forward.

func (*MultiHeadLatentAttention[T]) Forward added in v0.2.1 

func (m *MultiHeadLatentAttention[T]) Forward(ctx context.Context, inputs ...*tensor.TensorNumeric[T]) (*tensor.TensorNumeric[T], error)

Forward computes the MLA forward pass. Input: [batch, seqLen, hidden] Output: [batch, seqLen, hidden]

func (*MultiHeadLatentAttention[T]) OpType added in v0.2.1 

func (m *MultiHeadLatentAttention[T]) OpType() string

OpType returns the layer operation type.

func (*MultiHeadLatentAttention[T]) OutputShape added in v0.2.1 

func (m *MultiHeadLatentAttention[T]) OutputShape() []int

OutputShape returns the output shape.

func (*MultiHeadLatentAttention[T]) Parameters added in v0.2.1 

func (m *MultiHeadLatentAttention[T]) Parameters() []*graph.Parameter[T]

Parameters returns all trainable parameters.

func (*MultiHeadLatentAttention[T]) SetSaver added in v1.49.0 

func (m *MultiHeadLatentAttention[T]) SetSaver(sv graph.Saver[T])

SetSaver implements graph.SaverAware (ztensor ADR 006): the compressed KV latent and pre-wO attention output are consumed by Backward and expensive to recompute, so they are registered with the save-for-backward contract.

type NSACoarseCompression added in v1.28.0 

type NSACoarseCompression[T tensor.Numeric] struct {
	// contains filtered or unexported fields
}

NSACoarseCompression implements the coarse-grained token compression path of Native Sparse Attention (NSA). It divides the KV sequence into blocks of B tokens, computes block-level attention scores by averaging key representations per block, selects top-c blocks per query position, and attends to the selected blocks at full resolution.

func NewNSACoarseCompression added in v1.28.0 

func NewNSACoarseCompression[T tensor.Numeric](
	engine compute.Engine[T],
	ops numeric.Arithmetic[T],
	blockSize, topBlocks, numHeads, numKVHeads, headDim int,
) *NSACoarseCompression[T]

NewNSACoarseCompression creates a new NSACoarseCompression layer.

Parameters:

  • engine: compute engine for tensor operations
  • ops: arithmetic operations for the numeric type
  • blockSize: number of tokens per KV block (B)
  • topBlocks: number of blocks to select per query position (c)
  • numHeads: number of query attention heads
  • numKVHeads: number of key/value heads
  • headDim: dimension of each attention head

func (*NSACoarseCompression[T]) Attributes added in v1.28.0 

func (nsa *NSACoarseCompression[T]) Attributes() map[string]interface{}

Attributes returns the layer configuration.

func (*NSACoarseCompression[T]) Backward added in v1.28.0 

func (nsa *NSACoarseCompression[T]) Backward(_ context.Context, _ types.BackwardMode, dOut *tensor.TensorNumeric[T], _ ...*tensor.TensorNumeric[T]) ([]*tensor.TensorNumeric[T], error)

Backward implements the straight-through estimator for the block selection. In the forward pass we use hard top-k; the backward pass passes gradients through as if the selection were soft (identity on the selected paths).

func (*NSACoarseCompression[T]) Forward added in v1.28.0 

func (nsa *NSACoarseCompression[T]) Forward(_ context.Context, inputs ...*tensor.TensorNumeric[T]) (*tensor.TensorNumeric[T], error)

Forward computes the coarse-grained compression attention.

Inputs:

  • Q: [batch, numHeads, seqQ, headDim]
  • K: [batch, numKVHeads, seqKV, headDim]
  • V: [batch, numKVHeads, seqKV, headDim]

Returns output with shape [batch, numHeads, seqQ, headDim].

The algorithm:

  1. Reshape K into blocks of size B and compute block-level keys by averaging.
  2. Compute coarse attention scores: Q @ blockKeys^T.
  3. Select top-c blocks per query position.
  4. Gather full-resolution K,V from selected blocks.
  5. Compute fine-grained attention on selected blocks.

func (*NSACoarseCompression[T]) OpType added in v1.28.0 

func (nsa *NSACoarseCompression[T]) OpType() string

OpType returns the operation type identifier.

func (*NSACoarseCompression[T]) OutputShape added in v1.28.0 

func (nsa *NSACoarseCompression[T]) OutputShape() []int

OutputShape returns the output shape from the last forward call.

func (*NSACoarseCompression[T]) Parameters added in v1.28.0 

func (nsa *NSACoarseCompression[T]) Parameters() []*graph.Parameter[T]

Parameters returns nil (no trainable parameters).

type NSAFineSelection added in v1.28.0 

type NSAFineSelection[T tensor.Numeric] struct {
	// contains filtered or unexported fields
}

NSAFineSelection implements the fine-grained token selection path of Natively Sparse Attention (NSA). For each query position it scores all KV positions via Q*K^T / sqrt(d), selects the top-f tokens by score, and computes softmax attention only over those selected tokens.

func NewNSAFineSelection added in v1.28.0 

func NewNSAFineSelection[T tensor.Numeric](
	engine compute.Engine[T],
	topTokens, numHeads, numKVHeads, headDim int,
) *NSAFineSelection[T]

NewNSAFineSelection creates a new NSAFineSelection layer.

Parameters:

  • engine: compute engine for tensor operations
  • topTokens: number of tokens to select per query position (f)
  • numHeads: number of query heads
  • numKVHeads: number of key/value heads
  • headDim: dimension of each head

func (*NSAFineSelection[T]) Forward added in v1.28.0 

func (n *NSAFineSelection[T]) Forward(
	ctx context.Context,
	Q, K, V *tensor.TensorNumeric[T],
) (*tensor.TensorNumeric[T], error)

Forward computes the fine-grained token selection attention.

Inputs:

  • Q: [batch, numHeads, seqQ, headDim]
  • K: [batch, numKVHeads, seqKV, headDim]
  • V: [batch, numKVHeads, seqKV, headDim]

Returns:

  • output: [batch, numHeads, seqQ, headDim]

func (*NSAFineSelection[T]) SelectedIndices added in v1.28.0 

func (n *NSAFineSelection[T]) SelectedIndices(
	Q, K *tensor.TensorNumeric[T],
) []int

SelectedIndices computes and returns the top-f token indices per query position without computing the full attention output. This is useful for testing that the selection logic matches manual computation.

Returns indices with shape [batch, numHeads, seqQ, topTokens] (sorted descending by score).

type NSAWindowAttention added in v1.28.0 

type NSAWindowAttention[T tensor.Numeric] struct {
	// contains filtered or unexported fields
}

NSAWindowAttention implements the sliding window attention path for Native Sparse Attention (NSA). It takes pre-projected Q, K, V tensors with shape [batch, heads, seq, dim] and applies causal sliding window attention, restricting each query to attend only to keys within a fixed-size window.

func NewNSAWindowAttention added in v1.28.0 

func NewNSAWindowAttention[T tensor.Numeric](
	engine compute.Engine[T],
	ops numeric.Arithmetic[T],
	windowSize, numHeads, numKVHeads, headDim int,
) (*NSAWindowAttention[T], error)

NewNSAWindowAttention creates a new NSAWindowAttention layer.

Parameters:

  • engine: compute engine for tensor operations
  • ops: arithmetic operations for the numeric type
  • windowSize: number of past tokens each query can attend to (the window spans from max(0, q-windowSize) to q inclusive)
  • numHeads: number of query attention heads
  • numKVHeads: number of key/value attention heads
  • headDim: dimension of each attention head

func (*NSAWindowAttention[T]) Forward added in v1.28.0 

func (nw *NSAWindowAttention[T]) Forward(ctx context.Context, Q, K, V *tensor.TensorNumeric[T]) (*tensor.TensorNumeric[T], error)

Forward computes sliding window attention over pre-projected Q, K, V.

Input shapes:

  • Q: [batch, numHeads, seqQ, headDim]
  • K: [batch, numKVHeads, seqKV, headDim]
  • V: [batch, numKVHeads, seqKV, headDim]

Output shape: [batch, numHeads, seqQ, headDim]

For each query position q, attention is restricted to keys in the range [max(0, q - windowSize + 1), q] (causal sliding window). Positions outside this window receive -1e9 additive masking before softmax.

func (*NSAWindowAttention[T]) Scale added in v1.28.0 

func (nw *NSAWindowAttention[T]) Scale() float64

Scale returns the attention scaling factor (1/sqrt(headDim)).

func (*NSAWindowAttention[T]) WindowSize added in v1.28.0 

func (nw *NSAWindowAttention[T]) WindowSize() int

WindowSize returns the configured window size.

type NativeSparseAttention added in v1.28.0 

type NativeSparseAttention[T tensor.Numeric] struct {
	// contains filtered or unexported fields
}

NativeSparseAttention implements the full Native Sparse Attention (NSA) mechanism, combining three parallel attention paths:

  • Coarse: block-level compression attention
  • Fine: token-level selection attention
  • Window: sliding window local attention

The outputs are combined via learned per-head sigmoid gates: 

O = sigmoid(gateCoarse) * O_coarse + sigmoid(gateFine) * O_fine + sigmoid(gateWindow) * O_window

Gates are initialized to zero so that sigmoid(0) = 0.5 gives equal weighting at initialization.

func NewNativeSparseAttention added in v1.28.0 

func NewNativeSparseAttention[T tensor.Numeric](
	engine compute.Engine[T],
	ops numeric.Arithmetic[T],
	modelDim, numHeads, numKVHeads, blockSize, topBlocks, topTokens, windowSize int,
) (*NativeSparseAttention[T], error)

NewNativeSparseAttention creates a new NativeSparseAttention layer combining coarse, fine, and window attention paths with learned sigmoid gates.

Parameters:

  • engine: compute engine for tensor operations
  • ops: arithmetic operations for the numeric type
  • modelDim: model dimension (unused directly, reserved for projection layers)
  • numHeads: number of query attention heads
  • numKVHeads: number of key/value attention heads
  • blockSize: number of tokens per KV block for coarse path
  • topBlocks: number of blocks to select for coarse path
  • topTokens: number of tokens to select for fine path
  • windowSize: sliding window size for window path

func (*NativeSparseAttention[T]) Attributes added in v1.28.0 

func (nsa *NativeSparseAttention[T]) Attributes() map[string]interface{}

Attributes returns the layer configuration.

func (*NativeSparseAttention[T]) Backward added in v1.28.0 

func (nsa *NativeSparseAttention[T]) Backward(_ context.Context, _ types.BackwardMode, dOut *tensor.TensorNumeric[T], inputs ...*tensor.TensorNumeric[T]) ([]*tensor.TensorNumeric[T], error)

Backward computes gradients for the NativeSparseAttention layer. Gradients flow through the sigmoid gates to the gate parameters and through each attention path via straight-through estimation.

func (*NativeSparseAttention[T]) Forward added in v1.28.0 

func (nsa *NativeSparseAttention[T]) Forward(ctx context.Context, inputs ...*tensor.TensorNumeric[T]) (*tensor.TensorNumeric[T], error)

Forward computes NativeSparseAttention by running all three paths and combining with learned sigmoid gates.

Inputs:

  • Q: [batch, numHeads, seqQ, headDim]
  • K: [batch, numKVHeads, seqKV, headDim]
  • V: [batch, numKVHeads, seqKV, headDim]

Returns output with shape [batch, numHeads, seqQ, headDim].

func (*NativeSparseAttention[T]) OpType added in v1.28.0 

func (nsa *NativeSparseAttention[T]) OpType() string

OpType returns the operation type identifier.

func (*NativeSparseAttention[T]) OutputShape added in v1.28.0 

func (nsa *NativeSparseAttention[T]) OutputShape() []int

OutputShape returns the output shape from the last forward call.

func (*NativeSparseAttention[T]) Parameters added in v1.28.0 

func (nsa *NativeSparseAttention[T]) Parameters() []*graph.Parameter[T]

Parameters returns the trainable gate parameters.

type RopeScaler added in v0.2.0 

type RopeScaler[T tensor.Numeric] interface {
	ScaleRope(ctx context.Context, factor float64) error
}

RopeScaler is an interface for layers that support scaling of RoPE.

type ScaledDotProductAttention 

type ScaledDotProductAttention[T tensor.Numeric] struct {
	// contains filtered or unexported fields
}

ScaledDotProductAttention implements the scaled dot-product attention mechanism.

func NewBidirectionalSDPA added in v1.9.0 

func NewBidirectionalSDPA[T tensor.Numeric](engine compute.Engine[T], headDim int, opts ...ScaledDotProductAttentionOption[T]) *ScaledDotProductAttention[T]

NewBidirectionalSDPA creates a ScaledDotProductAttention layer with causal masking disabled. All positions attend to all other positions, which is the attention pattern used by encoder models such as BERT.

func NewScaledDotProductAttention 

func NewScaledDotProductAttention[T tensor.Numeric](engine compute.Engine[T], headDim int, opts ...ScaledDotProductAttentionOption[T]) *ScaledDotProductAttention[T]

NewScaledDotProductAttention creates a new ScaledDotProductAttention layer.

func (*ScaledDotProductAttention[T]) Backward 

func (sdpa *ScaledDotProductAttention[T]) Backward(ctx context.Context, mode types.BackwardMode, dOut, _, _, _ *tensor.TensorNumeric[T]) ([]*tensor.TensorNumeric[T], error)

Backward computes the gradients for ScaledDotProductAttention. dOut is the gradient from the subsequent layer.

func (*ScaledDotProductAttention[T]) Forward 

func (sdpa *ScaledDotProductAttention[T]) Forward(ctx context.Context, q, k, v, mask *tensor.TensorNumeric[T]) (*tensor.TensorNumeric[T], error)

Forward computes the scaled dot-product attention. Q, K, V are expected to be 3D tensors (batch_size, seq_len, head_dim). mask is an optional 4D tensor (batch_size, num_heads, seq_len_q, seq_len_k).

func (*ScaledDotProductAttention[T]) SetCausal added in v0.2.1 

func (sdpa *ScaledDotProductAttention[T]) SetCausal(causal bool)

SetCausal enables or disables causal (lower-triangular) masking.

func (*ScaledDotProductAttention[T]) SetSaver added in v1.49.0 

func (sdpa *ScaledDotProductAttention[T]) SetSaver(sv graph.Saver[T])

SetSaver implements graph.SaverAware. Owning nodes (AttentionHead, GroupedQueryAttention, ...) fan their Saver into SDPA so its cached tensors are attributed to the owning node and released after its Backward returns.

type ScaledDotProductAttentionOption added in v0.2.0 

type ScaledDotProductAttentionOption[T tensor.Numeric] func(*ScaledDotProductAttentionOptions[T])

ScaledDotProductAttentionOption applies an option to ScaledDotProductAttentionOptions.

func WithBidirectional added in v1.9.0 

func WithBidirectional[T tensor.Numeric]() ScaledDotProductAttentionOption[T]

WithBidirectional returns an option that disables causal masking, allowing every position to attend to every other position. This is required for encoder-style models such as BERT.

func WithHeadCounts added in v1.29.0 

func WithHeadCounts[T tensor.Numeric](numQueryHeads, numKVHeads int) ScaledDotProductAttentionOption[T]

WithHeadCounts sets the query and KV head counts, enabling the split-KV flash decode kernel for autoregressive decode with GQA support.

type ScaledDotProductAttentionOptions added in v0.2.0 

type ScaledDotProductAttentionOptions[T tensor.Numeric] struct {
	// contains filtered or unexported fields
}

ScaledDotProductAttentionOptions holds configuration options for ScaledDotProductAttention.

type SparseRoutedAttention added in v1.29.0 

type SparseRoutedAttention[T tensor.Numeric] struct {
	// contains filtered or unexported fields
}

SparseRoutedAttention routes queries to a subset of KV segments using cosine similarity scoring. It divides the KV sequence into fixed-size segments, computes cosine similarity between each query and segment centroids (mean key vectors), selects the top-k most similar segments, and performs scaled dot-product attention over the selected segments.

Position encoding uses document-wise RoPE so that position IDs reset at document boundaries during multi-document inference.

KV history is stored in a CompressedKVCache (via SparseRoutedKVCache interface) for efficient long-context inference.

func NewSparseRoutedAttention added in v1.29.0 

func NewSparseRoutedAttention[T tensor.Numeric](
	engine compute.Engine[T],
	ops numeric.Arithmetic[T],
	rope *embeddings.RotaryPositionalEmbedding[T],
	kvCache SparseRoutedKVCache[T],
	numHeads, numKVHeads, headDim, segmentSize, topK int,
) (*SparseRoutedAttention[T], error)

NewSparseRoutedAttention creates a new SparseRoutedAttention layer.

Parameters:

  • engine: compute engine for tensor operations (must support CosineSimilarity)
  • ops: arithmetic operations for the numeric type
  • rope: rotary positional embedding (supports document-wise mode)
  • kvCache: compressed KV cache for storing key-value pairs (satisfies SparseRoutedKVCache)
  • numHeads: number of query attention heads
  • numKVHeads: number of key/value attention heads
  • headDim: dimension of each attention head
  • segmentSize: number of tokens per KV segment for routing
  • topK: number of segments to select per query position

func (*SparseRoutedAttention[T]) Attributes added in v1.29.0 

func (sra *SparseRoutedAttention[T]) Attributes() map[string]interface{}

Attributes returns the layer configuration.

func (*SparseRoutedAttention[T]) Backward added in v1.29.0 

func (sra *SparseRoutedAttention[T]) Backward(_ context.Context, _ types.BackwardMode, dOut *tensor.TensorNumeric[T], _ ...*tensor.TensorNumeric[T]) ([]*tensor.TensorNumeric[T], error)

Backward computes gradients for the SparseRoutedAttention layer. Uses straight-through estimation for the routing selection.

func (*SparseRoutedAttention[T]) Forward added in v1.29.0 

func (sra *SparseRoutedAttention[T]) Forward(ctx context.Context, inputs ...*tensor.TensorNumeric[T]) (*tensor.TensorNumeric[T], error)

Forward computes sparse routed attention.

Inputs:

  • Q: [batch, numHeads, seqQ, headDim]
  • K: [batch, numKVHeads, seqKV, headDim]
  • V: [batch, numKVHeads, seqKV, headDim]

Returns output with shape [batch, numHeads, seqQ, headDim].

The algorithm:

  1. Divide K into segments of segmentSize tokens and compute segment centroids.
  2. For each query position, compute cosine similarity with all centroids.
  3. Select top-k segments per query.
  4. Gather full-resolution K,V from selected segments.
  5. Apply RoPE to Q and selected K.
  6. Compute scaled dot-product attention over selected segments.

func (*SparseRoutedAttention[T]) OpType added in v1.29.0 

func (sra *SparseRoutedAttention[T]) OpType() string

OpType returns the operation type identifier.

func (*SparseRoutedAttention[T]) OutputShape added in v1.29.0 

func (sra *SparseRoutedAttention[T]) OutputShape() []int

OutputShape returns the output shape from the last forward call.

func (*SparseRoutedAttention[T]) Parameters added in v1.29.0 

func (sra *SparseRoutedAttention[T]) Parameters() []*graph.Parameter[T]

Parameters returns nil (no trainable parameters).

func (*SparseRoutedAttention[T]) SetDocumentBoundaries added in v1.29.0 

func (sra *SparseRoutedAttention[T]) SetDocumentBoundaries(boundaries []int)

SetDocumentBoundaries sets document boundary positions for document-wise RoPE. Position IDs reset to 0 at each boundary so each document receives independent positional encoding. Pass nil to disable.

type SparseRoutedKVCache added in v1.29.0 

type SparseRoutedKVCache[T tensor.Numeric] interface {
	NumLayers() int
	SeqLen() int
	Reset()
}

SparseRoutedKVCache is the interface for compressed KV caches used by SparseRoutedAttention. This is satisfied by generate.CompressedKVCache and avoids an import cycle with the generate package.

Source Files

View all Source files

Jump to

Keyboard shortcuts

? : This menu
/ : Search site
f or F : Jump to
y or Y : Canonical URL
go.dev uses cookies from Google to deliver and enhance the quality of its services and to analyze traffic. Learn more.