Repository: ollama/ollama
File: model/models/olmo3/model.go (Go)

Code before:
		kvcache.NewCausalCache(m.Shift),
	)

	return &m, nil
}

type SelfAttention struct {
	Query  *nn.Linear  `gguf:"attn_q"`
	Key    *nn.Linear  `gguf:"attn_k"`
	Value  *nn.Linear  `gguf:"attn_v"`
	Output *nn.Linear  `gguf:"attn_output"`
	QNorm  *nn.RMSNorm `gguf:"attn_q_norm"`
	KNorm  *nn.RMSNorm `gguf:"attn_k_norm"`
}

func (o Options) applyRotaryPositionEmbeddings(ctx ml.Context, states, positions ml.Tensor, isSWA bool) ml.Tensor {
	freqScale := float32(1.0)
	ropeOpts := []func(*rope.Options){rope.WithTypeNeoX()}

	if !isSWA {
		freqScale = 1. / o.ropeScale
		if o.originalContextLength > 0 {
			ropeOpts = append(ropeOpts,
				rope.WithOriginalContextLength(o.originalContextLength),
				rope.WithExtrapolationFactor(o.ropeExtrapolation),
				rope.WithAttentionFactor(o.attnFactor),
			)
		}
	}

	return nn.RoPE(ctx, states, positions, o.hiddenSize/o.numHeads, o.ropeBase, freqScale, ropeOpts...)
}

func (sa *SelfAttention) Forward(ctx ml.Context, hiddenState, positions ml.Tensor, cache kvcache.Cache, m *Model, isSWA bool) ml.Tensor {
	batchSize := hiddenState.Dim(1)
	headDim := m.hiddenSize / m.numHeads

	query := sa.Query.Forward(ctx, hiddenState)
	query = sa.QNorm.Forward(ctx, query, m.eps)
	query = query.Reshape(ctx, headDim, m.numHeads, batchSize)
	query = m.Options.applyRotaryPositionEmbeddings(ctx, query, positions, isSWA)

	key := sa.Key.Forward(ctx, hiddenState)
	key = sa.KNorm.Forward(ctx, key, m.eps)
	key = key.Reshape(ctx, headDim, m.numKVHeads, batchSize)
	key = m.Options.applyRotaryPositionEmbeddings(ctx, key, positions, isSWA)

	value := sa.Value.Forward(ctx, hiddenState)
	value = value.Reshape(ctx, headDim, m.numKVHeads, batchSize)

	attention := nn.Attention(ctx, query, key, value, 1.0/math.Sqrt(float64(headDim)), cache)

Review comment (suggestion):
are you sure this is correct? ggml already sets attention factor to `1 + 0.1 * log(scale)` producing final value of `attention_factor * (1 + 0.1 * log(scale))`

Code after:
		kvcache.NewCausalCache(m.Shift),
	)

	return &m, nil
}

type SelfAttention struct {
	Query  *nn.Linear  `gguf:"attn_q"`
	Key    *nn.Linear  `gguf:"attn_k"`
	Value  *nn.Linear  `gguf:"attn_v"`
	Output *nn.Linear  `gguf:"attn_output"`
	QNorm  *nn.RMSNorm `gguf:"attn_q_norm"`
	KNorm  *nn.RMSNorm `gguf:"attn_k_norm"`
}

func (o Options) applyRotaryPositionEmbeddings(ctx ml.Context, states, positions ml.Tensor, isSWA bool) ml.Tensor {
	freqScale := float32(1.0)
	ropeOpts := []func(*rope.Options){rope.WithTypeNeoX()}

	if !isSWA {
		freqScale = 1. / o.ropeScale
		if o.originalContextLength > 0 {
			ropeOpts = append(ropeOpts,
				rope.WithOriginalContextLength(o.originalContextLength),
				rope.WithExtrapolationFactor(o.ropeExtrapolation),
			)
		}
	}

	return nn.RoPE(ctx, states, positions, o.hiddenSize/o.numHeads, o.ropeBase, freqScale, ropeOpts...)
}

func (sa *SelfAttention) Forward(ctx ml.Context, hiddenState, positions ml.Tensor, cache kvcache.Cache, m *Model, isSWA bool) ml.Tensor {
	batchSize := hiddenState.Dim(1)
	headDim := m.hiddenSize / m.numHeads

	query := sa.Query.Forward(ctx, hiddenState)
	query = sa.QNorm.Forward(ctx, query, m.eps)
	query = query.Reshape(ctx, headDim, m.numHeads, batchSize)
	query = m.Options.applyRotaryPositionEmbeddings(ctx, query, positions, isSWA)

	key := sa.Key.Forward(ctx, hiddenState)
	key = sa.KNorm.Forward(ctx, key, m.eps)
	key = key.Reshape(ctx, headDim, m.numKVHeads, batchSize)
	key = m.Options.applyRotaryPositionEmbeddings(ctx, key, positions, isSWA)

	value := sa.Value.Forward(ctx, hiddenState)
	value = value.Reshape(ctx, headDim, m.numKVHeads, batchSize)

	attention := nn.Attention(ctx, query, key, value, 1.0/math.Sqrt(float64(headDim)), cache)
	attention = attention.Reshape(ctx, m.hiddenSize, batchSize)