Livt.ML

src/attention/ScaledDotProductAttention.lvt

namespace Livt.ML.Attention

using Livt.Math.Arithmetic
using Livt.ML.Activation

/**
 * Computes scaled dot-product attention over a four-slot key/value cache.
 */
component ScaledDotProductAttention
{
	// Number of integer elements in query, key, value, and output vectors.
	public const VECTOR_SIZE: int = 8

	// Maximum number of cached key/value pairs.
	public const CACHE_SIZE: int = 4

	// Integer divisor applied to raw dot-product scores.
	public const SCORE_SCALE: int = 3

	// Fixed-point scale used by SoftmaxApprox weights.
	public const WEIGHT_SCALE: int = 256

	// Owned inference subcomponents sharing the caller clock/reset context.
	divider: IntegerDivision<32>
	kvcache: KVCache
	softmax: SoftmaxApprox

	// Latest computed result, retained for the caller.
	output: int[8]

	/**
	 * Owns cache, softmax, and scheduled score division; serialize operations.
	 */
	new()
	{
		divider = new IntegerDivision<32>()
		kvcache = new KVCache()
		softmax = new SoftmaxApprox()
	}

	/**
	 * Adds one key/value pair to the attention cache.
	 */
	public fn Write(key: int[8], val: int[8])
	{
		kvcache.Write(key, val)
	}

	/**
	 * Computes attention for query against the cached key/value pairs.
	 */
	public fn Compute(query: int[8])
	{
		var activeCount = kvcache.GetSize()

		if (activeCount == 0) {
			for (var elem = 0; elem < VECTOR_SIZE; elem++) {
				output[elem] = 0
			}

			return
		}

		var scores: int[8]

		for (var slot = 0; slot < activeCount; slot++) {
			var dot = 0

			for (var elem = 0; elem < VECTOR_SIZE; elem++) {
				dot = dot + query[elem] * kvcache.GetKeyElement(slot, elem)
			}

			scores[slot] = divider.Divide(dot, SCORE_SCALE)
		}

		softmax.ComputeCount(scores, activeCount)

		for (var elem = 0; elem < VECTOR_SIZE; elem++) {
			var weighted = 0

			for (var slot = 0; slot < activeCount; slot++) {
				weighted = weighted + softmax.GetOutput(slot) * kvcache.GetValueElement(slot, elem)
			}

			output[elem] = weighted / WEIGHT_SCALE
		}
	}

	/**
	 * Returns one attention output element, or zero when index is out of range.
	 */
	public fn GetOutput(index: int) int
	{
		if (index < 0) {
			return 0
		}

		if (index >= VECTOR_SIZE) {
			return 0
		}

		return output[index]
	}

	/**
	 * Returns the number of key/value pairs currently cached.
	 */
	public fn GetCacheSize() int
	{
		return kvcache.GetSize()
	}

	/**
	 * Clears the owned key/value cache.
	 */
	public fn ClearCache()
	{
		kvcache.Clear()
	}
}