Livt.ML

tests/IncrementalTransformerKernelsTest.lvt

namespace Livt.ML.Tests

using Livt.IO
using Livt.ML.Storage
using Livt.ML.Transformer

/**
 * Deterministic tensor-weight provider for transformer kernel tests; selection is local to
 * this shared fixture.
 */
component IncrementalWeights : IWeightTensor
{
	// Shared fixture instances, signals, and observation state.
	selected: int

	/**
	 * Selects a deterministic fixture tensor and reports whether its ID is supported.
	 */
	override fn Select(tensor: int) bool
	{
		selected = tensor
		return tensor >= 0 && tensor <= 2
	}

	/**
	 * Returns the row count of the selected fixture tensor; select a supported tensor first.
	 */
	override fn Rows() int
	{
		if (selected == 0) {
			return 4
		}

		if (selected == 1) {
			return 32128
		}

		return 3
	}

	/**
	 * Returns the column count of the selected fixture tensor.
	 */
	override fn Columns() int
	{
		return 1
	}

	/**
	 * Returns the deterministic weight for the selected fixture tensor and row-major index.
	 * Tests supply valid indices.
	 */
	override fn Read(index: int) int
	{
		if (selected == 0) {
			return 0
		}

		if (selected == 1) {
			if (index >= 32126) {
				return 512
			}

			return -256
		}

		if (index == 0) {
			return -256
		}

		if (index == 1) {
			return 16384
		}

		return 32767
	}

	/**
	 * Returns false because this fixture supplies only successful weight reads.
	 */
	override fn Failed() bool
	{
		return false
	}
}

/**
 * Exercises IncrementalTransformerKernels with a shared fixture instance. Each scenario
 * establishes the inputs and state it needs.
 */
@Test
component IncrementalTransformerKernelsTest
{
	// Shared fixture instances, signals, and observation state.
	ram: Ram<int, 128>
	scratch: BlockRam<int, 8>
	scratchAccess: IRam<int>
	weights: IncrementalWeights
	tables: SmallTransformerTables
	kernels: FixedTransformerKernels<4, 8, 16, 8, 8, 1, IncrementalWeights, SmallTransformerTables, Ram<int, 128>, BlockRam<int, 8>>

	/**
	 * Creates the shared fixture instances (Ram, IncrementalWeights, SmallTransformerTables,
	 * FixedTransformerKernels). Tests initialize the state required by their scenario.
	 */
	new()
	{
		ram = new Ram<int, 128>()
		scratch = new BlockRam<int, 8>()
		scratchAccess = scratch
		weights = new IncrementalWeights()
		tables = new SmallTransformerTables()
		kernels = new FixedTransformerKernels<4, 8, 16, 8, 8, 1, IncrementalWeights, SmallTransformerTables, Ram<int, 128>, BlockRam<int, 8>>(weights, tables, ram, scratch)
	}

	/**
	 * Checks the causal does not read future and cross masks case.
	 */
	@Test
	fn CausalDoesNotReadFutureAndCrossMasks()
	{
		ram.Write(0, 0)
		ram.Write(8, 0)
		ram.Write(9, 2147483647)
		ram.Write(10, 2147483647)
		ram.Write(16, 768)
		ram.Write(17, 2147483647)
		ram.Write(18, 2147483647)

		for (var i = 0; i < 3; i++) {
			ram.Write(32 + i, 1)
		}

		for (var i = 0; i < 7; i++) {
			ram.Write(40 + i, 0)
		}

		assert kernels.CachedAttention(0, 0, 8, 16, 24, 32, 3, 1, 1, 0, true, 40)
		assert ram.Read(24) == 768
		ram.Write(9, 0)
		ram.Write(17, 256)
		ram.Write(34, 0)
		assert kernels.CachedAttention(-1, 0, 8, 16, 24, 32, 3, 1, 1, 0, false, -1)
		assert ram.Read(24) == 512
		ram.Write(32, 0)
		ram.Write(33, 0)
		assert !kernels.CachedAttention(-1, 0, 8, 16, 24, 32, 3, 1, 1, 0, false, -1)
		assert ram.Read(24) == 512
		assert !kernels.CachedAttention(0, 0, 8, 16, 24, 32, 3, 1, 1, 4, true, 40)
	}

	/**
	 * Checks the entire vocabulary lowest tie and clipped logits case.
	 */
	@Test
	fn EntireVocabularyLowestTieAndClippedLogits()
	{
		ram.Write(0, 256)
		assert kernels.ProjectArgMax(1, 0, 32128, 1) == 32126
		ram.Write(0, 32767)
		assert kernels.ProjectArgMax(2, 0, 3, 1) == 1
		assert kernels.ProjectArgMax(2, 0, 4, 1) == -1
		assert kernels.ProjectArgMax(2, 128, 3, 1) == -1
	}

	/**
	 * Shares full-capacity scratch across encoder/cached attention, failures and new masks.
	 * Distinct scores select the last valid value; poison detects stale array contents.
	 */
	@Test
	fn SharedScratchReuseAfterFailure()
	{
		for (var i = 0; i < 8; i++) {
			scratchAccess.Write(i, 2147483647)
		}

		for (var i = 0; i < 4; i++) {
			ram.Write(i, 256)
			ram.Write(8 + i, i * 256)
			ram.Write(16 + i, (i + 1) * 256)
			ram.Write(32 + i, 1)
		}

		assert kernels.Attention(0, 0, 8, 16, 24, 32, 4, 1, 1)

		for (var i = 0; i < 4; i++) {
			assert ram.Read(24 + i) == 1024
		}

		ram.Write(11, 32768)
		assert !kernels.CachedAttention(-1, 0, 8, 16, 24, 32, 4, 1, 1, 0, false, -1)
		ram.Write(35, 0)
		assert kernels.CachedAttention(-1, 0, 8, 16, 24, 32, 4, 1, 1, 0, false, -1)
		assert ram.Read(24) == 768
		ram.Write(11, 768)

		for (var i = 1; i < 4; i++) {
			ram.Write(32 + i, 0)
		}

		assert kernels.Attention(0, 0, 8, 16, 24, 32, 4, 1, 1)

		for (var i = 0; i < 4; i++) {
			assert ram.Read(24 + i) == 256
		}
	}

	/**
	 * Checks projection clipping, lowest-index ties, invalid spans and invalid activation reuse.
	 */
	@Test
	fn SmallProjectionBoundsAndClippedTie()
	{
		ram.Write(0, 32767)
		assert kernels.ProjectArgMax(2, 0, 3, 1) == 1
		assert kernels.ProjectArgMax(2, 0, 4, 1) == -1
		assert kernels.ProjectArgMax(2, 128, 3, 1) == -1
		ram.Write(0, 32768)
		assert kernels.ProjectArgMax(2, 0, 3, 1) == -1
		ram.Write(0, 32767)
		assert kernels.ProjectArgMax(2, 0, 3, 1) == 1
	}

}