WebAssembly SIMD for Exact All-Pairs Clustering: From 428 s to 8 s in Node.js

A clustering job in a Node.js and TypeScript service compares every pair of article embeddings and joins two articles when their cosine similarity reaches 0.65. At 30,000 articles it took 428 s. After moving the inner loop into a hand-written WebAssembly SIMD kernel and running it on every core, the same job takes about 8 s, and it produces exactly the same clusters.

The job is still quadratic. This post is about what a constant factor of roughly 50 is worth, what two clever attempts at removing work did not do, and how to keep an optimisation exact when floating point is involved.

The problem: exact all-pairs cosine similarity

Every article has an embedding, and the embeddings are normalised to unit length, so cosine similarity is a dot product. Clustering here means: for each pair of articles, compute the dot product, and if it reaches the threshold, the pair is joined. With n articles that is n * (n - 1) / 2 dot products, each over the whole vector.

The original implementation was the obvious one, a double loop in plain TypeScript with number arithmetic. It is exact and easy to read, and it scales the way its shape says. Measured on an Apple Silicon desktop, with the inputs built from a real 48-hour window scaled up by seeded noise and shifted times:

articles	old loop	new kernel
6,000	16.6 s	0.4 s
12,000	68.8 s	1.4 s
30,000	428.3 s	7.6 to 8.3 s
60,000	not run	32.9 s

Doubling the input roughly quadruples the old time, as expected. The new column grows the same way, only from a much smaller starting point. Nothing about the algorithm changed.

What did not work: removing pairs exactly

The first instinct with a quadratic job is to skip pairs that cannot pass. It has to be exact, so the pruning has to be provably safe, not a heuristic.

PCA plus Cauchy-Schwarz. Project the vectors onto their leading principal components. The dot product splits into a part in that subspace and a remainder, and Cauchy-Schwarz bounds the remainder by the product of the two residual norms. If the subspace part plus that bound is still under 0.65, the pair is safely skipped. The bound is valid, so the result stays exact. It removed only 28% of the pairs, and the run took 15.7 s against 19.5 s for the plain loop. Two reasons: random pairs of these embeddings already sit near cosine 0.3, so the bound has little room to cut below 0.65, and the vectors are not low-rank, so the residual norms stay large.

Triangle-inequality pivots. Pick a few pivot vectors, precompute every article's distance to them, and skip pairs whose pivot distances differ by too much. This failed for the same reason. When most pairs are already at a similar middle distance, distances to a pivot barely separate them.

The lesson is general. Exact pruning needs structure in the data, such as low rank or clusters far apart, and embedding vectors in a narrow, busy neighbourhood do not have it. When the work cannot be removed, make each unit of work cheaper and spread it out.

What worked: a SIMD kernel on every core

Two multipliers stack. A dot product is a long run of multiply-and-add over independent lanes, and WebAssembly has a 128-bit vector type for exactly that. And the pairs are independent of each other, so they split across worker_threads.

flowchart LR
  A["a: 4 floats"] --> M["f32x4.mul"]
  B["b: 4 floats"] --> M
  M --> S["f32x4.add into 4-lane accumulator"]
  S --> L{"more vectors?"}
  L -- yes --> A
  L -- no --> R["add the 4 lanes into one float"]

Platform support

Fixed-width SIMD in WebAssembly is a 128-bit packed proposal, described by the proposal text and the V8 write-up as standardising fixed-width 128-bit operations on a v128 type. caniuse lists Chrome 91, Edge 91, Firefox 89 and Safari 16.4 as the first versions with it. For Node.js, the 16.4.0 release notes show the V8 update to 9.1.269.36, and WebAssembly.validate on a SIMD module returns true on Node 24 without any flag.

Relaxed SIMD is a separate proposal. It adds instructions whose results may vary by hardware, such as a fused multiply-add that rounds once or twice depending on the machine. That is the opposite of what an exact job wants, so this kernel uses only the fixed-width, deterministic instructions.

The kernel

A dot product over a length that is a multiple of 4, in WebAssembly text. Each iteration loads four floats from each vector, multiplies them lane by lane, and adds into a four-lane accumulator. The four lanes are summed once at the end.

(module
  (memory (export "mem") 1)
  (func (export "dot") (param $a i32) (param $b i32) (param $n i32) (result f32)
    (local $acc v128) (local $end i32)
    (local.set $end (i32.add (local.get $a) (i32.shl (local.get $n) (i32.const 2))))
    (local.set $acc (v128.const f32x4 0 0 0 0))
    (block $done
      (loop $next
        (br_if $done (i32.ge_u (local.get $a) (local.get $end)))
        (local.set $acc
          (f32x4.add (local.get $acc)
            (f32x4.mul (v128.load (local.get $a)) (v128.load (local.get $b)))))
        (local.set $a (i32.add (local.get $a) (i32.const 16)))
        (local.set $b (i32.add (local.get $b) (i32.const 16)))
        (br $next)))
    (f32.add
      (f32.add (f32x4.extract_lane 0 (local.get $acc)) (f32x4.extract_lane 1 (local.get $acc)))
      (f32.add (f32x4.extract_lane 2 (local.get $acc)) (f32x4.extract_lane 3 (local.get $acc))))))

a and b are byte offsets into the module's memory and n counts floats. Embedding lengths are usually a multiple of 4 already, and padding with zeros fixes the rest without changing any dot product. To assemble it without a toolchain install, the wabt npm package works from a script:

import wabt from 'wabt';
import fs from 'node:fs';

const w = await wabt();
const module = w.parseWat('dot.wat', fs.readFileSync('dot.wat', 'utf8'), { simd: true });
fs.writeFileSync('dot.wasm', module.toBinary({}).buffer);

Splitting the pairs across cores

Each worker handles a strided set of rows, so every worker gets a similar number of pairs; the first rows have the most partners and would overload one worker if rows were cut into contiguous blocks. The embeddings live in a SharedArrayBuffer, so the main thread does not copy them per worker. Inside a worker, the matrix is copied once into that worker's WebAssembly memory, because the kernel reads from its own linear memory.

const { instance } = await WebAssembly.instantiate(wasmBytes);
const { mem, dot } = instance.exports;
mem.grow(Math.ceil((n * DIM * 4) / 65536));
new Float32Array(mem.buffer, 0, n * DIM).set(new Float32Array(sab));

for (let i = k; i < n; i += cores) {
  for (let j = i + 1; j < n; j++) {
    const s = dot(i * DIM * 4, j * DIM * 4, DIM);
    // ...compare s with the threshold, see the next section
  }
}

Here k is the worker's index and cores the worker count, and each worker posts back the pairs it accepted. A real job needs more care with memory, but the shape is this.

How exactness was kept

The old loop used double precision. The kernel uses single precision, in a different summation order. Two honest consequences follow: the f32 result can differ from the double result in the last digits, and a pair sitting right at 0.65 can land on the other side of the line.

The fix is a guard band. Any pair whose f32 score is within about 0.0005 of the threshold is not decided by the kernel at all. It is recomputed in double, in the same loop order as the old code, and that value decides. A pair outside the band is far enough from the line that the f32 error cannot cross it. So every join decision equals the one the old code would have made.

const MARGIN = 0.0005;

function exactDot(v, i, j) {
  let s = 0;
  for (let d = 0; d < DIM; d++) s += v[i * DIM + d] * v[j * DIM + d];
  return s;
}

// inside the pair loop
if (Math.abs(s - THRESHOLD) <= MARGIN) {
  if (exactDot(v, i, j) >= THRESHOLD) found.push(i * n + j);
} else if (s > THRESHOLD) {
  found.push(i * n + j);
}

The band is a tolerance on f32 error, so it has to be wider than the error and as narrow as possible, because every pair inside it costs a slow recomputation.

How the equivalence was tested

"Identical output" was checked three ways, and the order matters:

• A kernel test that builds vectors whose scores fall around the threshold and compares the decision against the double loop.
• A property test over 13 seeded random windows, comparing clusters, members and representatives against the old implementation.
• A replay on a real 48-hour window of 6,574 articles: identical clusters, members and representatives, 19.96 s with the old code and 0.79 s with the new.

A note on the snippets above: they are a standalone illustration of the technique, not the production code. Run together as a small program on seeded vectors, they accept the same pairs as the plain double loop. That is a smoke test and says nothing about the margin's adequacy on its own. The property test and the kernel test, which target the band directly, are what carry the claim.

The fallback

The fast path is an optimisation, never a requirement. Under 800 items the setup cost of workers and memory is not worth paying, so the old loop runs. It also runs when SIMD or workers are not available, and when memory or candidate limits would be exceeded. The old loop is slower and exactly as correct, so a failure in the fast path degrades the job's speed and never its result.

What it costs

• It is still quadratic. The constant is about 50 times smaller, so the curve is the same shape with a late knee; at 60,000 articles the run is 32.9 s.
• It uses every core for a few seconds. On a shared machine that is a real trade: the job finishes quickly, and everything else on that machine waits during that window.
• It is more code to own: a hand-written module, a worker protocol and a guard band, in place of one readable loop. The tests above exist because of that.
• Single precision is the reason the guard band exists at all.

When to reach for SIMD

• When the hot loop is arithmetic over contiguous numbers and the profile says that loop is the time. Dot products, distances, checksums and image kernels are the usual cases.
• After cheaper wins are gone. Exact pruning is worth a try, and here it taught what the data looks like, but it only helps when the data has structure to exploit.
• When the work is independent, so that cores multiply the gain.

To keep an optimisation exact, treat it as a candidate generator and let the old code arbitrate the cases that matter. Keep the old implementation as the fallback, put a guard band around every decision boundary, and test equality of the outputs, not just similarity of the timings.