This is an implementation walkthrough: it makes the shape of serving an attention computation over HTTP concrete in Rust, and it corrects a naming choice and a softmax bug found in the original code. The Rust code below is a pedagogical scaffold that wires scaled dot-product attention [1] into an HTTP service — named precisely as such, not as the IO-aware FlashAttention algorithm of Dao et al. [2], whose tiled kernel is the production target this scaffold points to (§4). The deliverable is the wiring and the naming precision, not a benchmark: no performance is measured because none is what this example is for.

Abstract

This note revises an earlier draft of the same Rust project ("flash-attention mechanizer, HOF cognition, distributed representation management") to be precisely scoped: a small, readable example of exposing an attention-style computation as a service in Rust — not a claim to FlashAttention itself. Three scope corrections carry the revision: (a) the kernel is named for what it is — standard scaled dot-product attention [1], with FlashAttention [2] cited as the production target it is not; (b) a softmax bug in the original snippet is identified and fixed, shown in §4; and (c) the deliverable is the wiring and the composition, not throughput, so no performance claim is made because none was needed to demonstrate what this example demonstrates. Rust is discussed as an implementation language; the design is described in Input–Process–Output terms.

1. Introduction

Making an attention computation available behind an HTTP endpoint is a reasonable, ordinary engineering exercise, and Rust is a sound choice for it. This revision keeps that walkthrough and sharpens one thing: naming. The module — renamed here from flash_attention to attention — implements standard scaled dot-product attention, not the tiled, IO-aware exact-attention algorithm that the name FlashAttention refers to in the literature [2]; that algorithm is the production target, cited in §4 as where this scaffold's naive score-matrix approach gets replaced. Naming it precisely, rather than under the FlashAttention name alongside "optimal performance," keeps the example honest about what it demonstrates: the wiring, not a measured result.

2. Related Work

Attention. The kernel computes (a sketch of) scaled dot-product attention [1]. FlashAttention. The real IO-aware exact-attention algorithm [2] tiles the computation to avoid materializing the n × n score matrix in high-bandwidth memory; the code here does no such tiling and is not that algorithm. Composition. Structuring the service as small, reusable stages (project → score → aggregate → serve) follows the ordinary functional-composition rationale for modularity [3, 4, 5]; it is a code-organization choice, not a novel mechanism.

3. Design (IPO)

4. The kernel — and a correctness caveat

The original attention kernel was:

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();
    // NOTE: this is NOT a correct softmax. It exponentiates the scores,
    // sums per row into a vector of reciprocals, and multiplies that by V.
    // A correct softmax attention divides exp(scores) elementwise by its
    // row-sum BEFORE multiplying by V.
    let weights = scores
        .mapv(f32::exp)
        .sum_axis(Axis(1))
        .mapv(|sum| 1.0 / sum);
    weights.dot(&v)
}

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

// Corrected sketch (still naive — materializes the full score matrix,
// i.e. the opposite of what FlashAttention [2] avoids):
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)); // broadcast divide
weights.dot(&v)

Even corrected, this materializes the full n × n score matrix and is therefore the naive memory profile [1] — precisely the profile FlashAttention exists to avoid [2]. So the module should not carry the FlashAttention name.

5. Serving it

The remaining project structure — a HOF wrapper that calls the kernel, a DistributedRepresenter that stores/loads representations to JSON, and a hyper-based server exposing /process and /retrieve — is straightforward, correct web plumbing and is retained from the original as a working example. The composition into small modules (flash_attention → rename to attention; hof_cognition; distributed_representation; server) is a reasonable code-organization choice [3, 4]. 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.

The wrapper and the route it serves, illustrative pseudocode, not a benchmarked implementation (elided: imports, error handling, the DistributedRepresenter body — see §3 for what each stage does):

pub struct HOF { attention: AttentionKernel }

impl HOF {
    pub fn process(&self, input: &Array2<f32>) -> Array2<f32> {
        self.attention.compute_attention(input) // the kernel in §4
    }
}

// POST /process: JSON floats in, attention output out
(&Method::POST, "/process") => {
    let input = Array2::from_shape_vec((1, input_vec.len()), input_vec)?;
    let output = hof.process(&input);
    Ok(Response::new(Body::from(json!(output).to_string())))
}

What this reference shows

This is a readable Rust reference for the shape of serving an attention computation over HTTP: Input → Process → Output, with every stage on the page — kernel, HOF wrapper, representer, HTTP server — so the composition reads as one legible pipeline. Each production concern is named exactly where it is solved.

References

  1. Ashish Vaswani et al. (2017). Attention Is All You Need. Advances in Neural Information Processing Systems (NeurIPS). arXiv:1706.03762.
  2. 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.
  3. 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]
  4. John Hughes (1989). Why Functional Programming Matters. The Computer Journal.
  5. Christopher Strachey (2000). Fundamental Concepts in Programming Languages. Higher-Order and Symbolic Computation. [Reprint of 1967 lecture notes]