Abstract

An attention computation exposed behind an HTTP endpoint is a systems object, not just a kernel. This note mechanizes scaled dot-product attention [1] as a composable higher-order pipeline in Rust — project → score → aggregate → serve — where each stage is a small, reusable function in the functional-programming tradition [2, 3, 4]. The kernel computes exact scaled dot-product attention [1]; the stage boundary around it is the design's point — it is clean enough that the tiled, IO-aware kernel of FlashAttention [5] is a drop-in swap for the scoring stage without touching the service that wraps it. Rust is the implementation language; the architecture is stated in Input–Process–Output terms.

1. Introduction

Making an attention computation available behind an HTTP endpoint is ordinary, load-bearing engineering, and Rust is a sound choice for it: a typed, memory-safe kernel served with no garbage-collector pauses. The design factors the service into small reusable stages — project the input to query/key/value, score, normalize, aggregate, serve — each a function composed into the next. The kernel implements exact scaled dot-product attention [1]. FlashAttention [5] is the tiled, IO-aware exact-attention algorithm that removes the n × n score matrix from high-bandwidth memory; the composition here draws the stage boundary precisely so that kernel drops into the scoring stage without disturbing the service around it. That swap boundary is the contribution.

2. Related Work

Attention. The kernel computes scaled dot-product attention [1] — project to Q/K/V, score Q·Kᵀ, scale, softmax, and aggregate against V.

FlashAttention. The IO-aware exact-attention algorithm [5] tiles the computation to avoid materializing the n × n score matrix in high-bandwidth memory. It is the kernel the scoring stage is built to accept: the service exposes a stage interface, and a tiled kernel satisfies it in place of the dense one.

Composition. Structuring the service as small, reusable stages (project → score → aggregate → serve) is the functional-composition rationale for modularity [2, 3, 4]. The composition is what makes the kernel a swappable component rather than a fused monolith.

3. Design (IPO)

4. The Kernel

Exact scaled dot-product attention [1] normalizes exp(scores) by its row-wise sum to a row-stochastic weight matrix A, then returns A · V:

pub fn compute_attention(&self, x: &Array2<f32>) -> Array2<f32> {
 let q = x.dot(&self.query);
 let k = x.dot(&self.key);
 let v = x.dot(&self.value);

 let scores = q.dot(&k.t()) / (x.shape()[1] as f32).sqrt();
 let exp = scores.mapv(f32::exp);
 let row_sums = exp.sum_axis(Axis(1));
 let weights = &exp / &row_sums.insert_axis(Axis(1)); // row-stochastic A
 weights.dot(&v)
}

This dense kernel materializes the full n × n score matrix — the standard memory profile of exact attention [1]. FlashAttention [5] is the tiled kernel that computes the same exact result without ever holding that matrix in high-bandwidth memory. The stage boundary in §3 is the reason the two are interchangeable here: the service consumes an attention output, not a specific kernel, so the memory profile is a property of the kernel placed in the scoring stage, chosen against [5] as the reference.

5. Serving It

The project structure is a HOF wrapper that calls the kernel, a DistributedRepresenter that stores and loads representations to JSON, and a hyper-based server exposing /process and /retrieve. Each is a stage: attention (the kernel of §4), hof_cognition (the composing wrapper), distributed_representation (the store), and server (transport). The composition into small modules is the functional-organization choice [2, 3] that keeps the kernel a swappable unit. [ILLUSTRATIVE] wiring: the example runs on 127.0.0.1:3000; curl -X POST .../process -d '[0.5,0.1,0.4,0.8]' exercises the kernel and curl .../retrieve returns a stored representation.

Evidence & Scope

The contribution is the composition and the swap boundary, not a throughput number. The kernel of §4 is exact scaled dot-product attention [1] and materializes the full score matrix; FlashAttention [5] is the tiled kernel the scoring stage is designed to accept in its place, and it is the reference the memory profile is measured against. The served example is illustrative wiring — single-row inputs, an in-memory representation store — demonstrating the stage interfaces end to end. The architectural claim stands on the boundary: a service that consumes an attention output rather than a fixed kernel makes the kernel a pure swap.

References

  1. Ashish Vaswani et al. (2017). Attention Is All You Need. Advances in Neural Information Processing Systems (NeurIPS). arXiv:1706.03762.
  2. John Backus (1978). Can Programming Be Liberated from the von Neumann Style? A Functional Style and Its Algebra of Programs. Communications of the ACM. [1977 ACM Turing Award Lecture]
  3. John Hughes (1989). Why Functional Programming Matters. The Computer Journal.
  4. Christopher Strachey (2000). Fundamental Concepts in Programming Languages. Higher-Order and Symbolic Computation. [Reprint of 1967 lecture notes]
  5. Tri Dao et al. (2022). FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. Advances in Neural Information Processing Systems (NeurIPS). arXiv:2205.14135.