Livt.ML

src/attention/CausalAttentionHead4Context8.lvt

namespace Livt.ML.Attention

using Livt.IO
using Livt.ML.Activation

/**
 * Reference causal scaled-dot-product attention for one GPT-Neo head.
 *
 * The fixed head width is four and context is eight. Dot products are divided
 * by sqrt(4)=2 exactly. Q1.15 softmax weights are applied to stored integer
 * values. Sixteen heads may reuse this component sequentially.
 */
component CausalAttentionHead4Context8
{
	// Fixed attention geometry.
	public const CONTEXT_LENGTH: int = 8
	public const HEAD_WIDTH: int = 4

	// Independent asynchronous vector stores and score normalization.
	queries: AsynchronousDistributedRam32x32
	keys: AsynchronousDistributedRam32x32
	values: AsynchronousDistributedRam32x32
	output: AsynchronousDistributedRam32x32
	softmax: StableSoftmax8

	// Live RAM boundary; scheduled methods own commands, reads stay combinational.
	queriesAddress: logic[5]
	queriesWriteEnable: bool
	queriesWriteData: logic[32]
	queriesReadData: logic[32]
	keysAddress: logic[5]
	keysWriteEnable: bool
	keysWriteData: logic[32]
	keysReadData: logic[32]
	valuesAddress: logic[5]
	valuesWriteEnable: bool
	valuesWriteData: logic[32]
	valuesReadData: logic[32]
	outputAddress: logic[5]
	outputWriteEnable: bool
	outputWriteData: logic[32]
	outputReadData: logic[32]

	/**
	 * Creates independent stores; SetElement initializes cells before Compute.
	 */
	new()
	{
		this.queries = new AsynchronousDistributedRam32x32()
		this.keys = new AsynchronousDistributedRam32x32()
		this.values = new AsynchronousDistributedRam32x32()
		this.output = new AsynchronousDistributedRam32x32()
		this.softmax = new StableSoftmax8()
	}

	/**
	 * Stores one query, key, and value element.
	 */
	public fn SetElement(token: int, element: int, query: int, key: int, value: int)
	{
		if (token < 0 || token >= CONTEXT_LENGTH) {
			return
		}

		if (element < 0 || element >= HEAD_WIDTH) {
			return
		}

		var address = token * HEAD_WIDTH + element

		state {
			queriesAddress = address as logic[5]
			queriesWriteData = query as logic[32]
			queriesWriteEnable = true
			keysAddress = address as logic[5]
			keysWriteData = key as logic[32]
			keysWriteEnable = true
			valuesAddress = address as logic[5]
			valuesWriteData = value as logic[32]
			valuesWriteEnable = true
		}

		state {
			queriesWriteEnable = false
			keysWriteEnable = false
			valuesWriteEnable = false
		}

		state {
		}
	}

	/**
	 * Computes every causal output using scoreStep units per natural-log unit.
	 */
	public fn Compute(scoreStep: int)
	{
		for (var token = 0; token < CONTEXT_LENGTH; token++) {
			var scores: int[8]

			for (var source = 0; source <= token; source++) {
				var dot = 0

				for (var element = 0; element < HEAD_WIDTH; element++) {
					state {
						queriesAddress = (token * HEAD_WIDTH + element) as logic[5]
						keysAddress = (source * HEAD_WIDTH + element) as logic[5]
					}

					state {
					}

					var queryValue = queriesReadData as int
					var keyValue = keysReadData as int
					var dotProduct = queryValue * keyValue

					state {
					}

					dot = dot + dotProduct
				}

				scores[source] = dot / 2
			}

			softmax.ComputeCount(scores, token + 1, scoreStep)

			for (var element = 0; element < HEAD_WIDTH; element++) {
				var weighted = 0

				for (var source = 0; source <= token; source++) {
					state {
						valuesAddress = (source * HEAD_WIDTH + element) as logic[5]
					}

					state {
					}

					var value = valuesReadData as int
					var weightedProduct = softmax.GetOutput(source) * value

					state {
					}

					weighted = weighted + weightedProduct
				}

				var result = weighted / StableSoftmax8.OUTPUT_SCALE

				state {
					outputAddress = (token * HEAD_WIDTH + element) as logic[5]
					outputWriteData = result as logic[32]
					outputWriteEnable = true
				}

				state {
					outputWriteEnable = false
				}

				state {
				}
			}
		}
	}

	/**
	 * Returns one output element, or zero for invalid coordinates.
	 */
	public fn GetOutput(token: int, element: int) int
	{
		if (token < 0 || token >= CONTEXT_LENGTH) {
			return 0
		}

		if (element < 0 || element >= HEAD_WIDTH) {
			return 0
		}

		state {
			outputAddress = (token * HEAD_WIDTH + element) as logic[5]
		}

		state {
		}

		return outputReadData as int
	}

	/**
	 * Connects inherited RAM ports without implicit scheduled field observations.
	 */
	process MemoryPorts[]()
	{
		this.queries.address = queriesAddress
		this.queries.writeEnable = queriesWriteEnable
		this.queries.writeData = queriesWriteData
		queriesReadData = this.queries.readData
		this.keys.address = keysAddress
		this.keys.writeEnable = keysWriteEnable
		this.keys.writeData = keysWriteData
		keysReadData = this.keys.readData
		this.values.address = valuesAddress
		this.values.writeEnable = valuesWriteEnable
		this.values.writeData = valuesWriteData
		valuesReadData = this.values.readData
		this.output.address = outputAddress
		this.output.writeEnable = outputWriteEnable
		this.output.writeData = outputWriteData
		outputReadData = this.output.readData
	}
}