The Engine

The experiment platform is a small, fast transformer inference engine in pure Java 21+ (re-implemented in this repository at code/attention/src/), built for one purpose: experimenting with attention machinery. Everything except attention is fixed — Llama-2 architecture, fp32, single sequence — and attention is a plugin. It runs Karpathy's tinyllamas checkpoints in llama2.c format; the default stories15M (15M parameters) loads in a fraction of a second and generates coherent English fast enough for tight experiment loops, with zero GC in the token loop.

The plugin seam

An experiment implements one interface:

public interface Attention {
  void init(Config config, Parallel pool);  // allocate all state here
  void reset();                             // new sequence
  void attend(int layer, int pos, float[] q, float[] k, float[] v, float[] out);
}

The engine hands over the post-RoPE query and this position's key/value; the implementation owns everything downstream — including the KV Cache. That is deliberate: experiments can change the cache representation itself (windowing, compression, recurrent state, no cache at all), not just the scoring rule. An implementation can also raise maxPositions to declare it isn't bound by a dense cache, letting generation run past the trained context with RoPE extrapolating.

Built-ins: causal (the exact reference), sliding:N, sink:s,w (the StreamingLLM-style sinks+window eviction baseline from Related Work), plus the experimental lattice, kmeans, and tiered memories.

The correctness harness

--diff teacher-forces the same tokens through two implementations sharing one set of weights and reports per-position logit divergence. Because both implementations use the same kernels, an exactly-equivalent formulation must produce a max |Δlogit| of literally zero — a sharp test when refactoring attention math. The build re-runs this harness on the standard probe:

Design notes

• Vectorized kernels via the JDK Vector API: dot products with unrolled FMA accumulators, softmax, saxpy, SwiGLU, and a custom vectorized exp (on macOS VectorOperators.EXP has no intrinsic, and the sampler's 32k-wide softmax was a large fraction of token latency before the fix).
• Zero-allocation hot path: every buffer is allocated at init and reused; the sampler ports llama2.c's xorshift RNG and top-p sampling exactly, which is what makes seeded stream runs bit-reproducible for verification.
• Threading: a fixed pool of spinning workers with static slices and one volatile epoch counter — no queues, no futures, no per-call allocation. Each output element is computed by a single thread, so results are deterministic regardless of thread count.
• RoPE stays in the engine (precomputed tables), since it's part of what the checkpoint's weights were trained against; implementations that want to own positional encoding un-rotate behind the interface.

Speed figures (tok/s) quoted on other pages are attested measurements, not build-verified numbers — see Verification Methodology.