gpu.test.mjs

import test from 'node:test';
import assert from 'node:assert/strict';
import * as g from './gpu.mjs';

const close = (actual, expected, tolerance, message) =>
  assert.ok(
    Math.abs(actual - expected) <= tolerance,
    `${message ?? ''} expected ~${expected}, got ${actual} (tolerance ${tolerance})`,
  );

/* ------------------------------------------------------------------ *
 * The spec table has to agree with itself.
 *
 * This is the test that matters most: if someone drops an invented number
 * into GPUS, the structural identity SMs x lanes x 2 x clock stops
 * reproducing the published TFLOPS and this fails.
 * ------------------------------------------------------------------ */

test('published FP32 TFLOPS is reproduced by SMs x lanes x 2 FLOP x clock', () => {
  for (const id of g.GPU_IDS) {
    const gpu = g.GPUS[id];
    const derived = g.derivedFp32Tflops(gpu);
    const relativeError = Math.abs(derived - gpu.fp32Tflops) / gpu.fp32Tflops;
    assert.ok(
      relativeError < 0.005,
      `${gpu.name}: derived ${derived.toFixed(2)} vs published ${gpu.fp32Tflops} TFLOPS (${(relativeError * 100).toFixed(2)}% off)`,
    );
  }
});

test('every GPU entry cites a source and has plausible structural fields', () => {
  for (const id of g.GPU_IDS) {
    const gpu = g.GPUS[id];
    assert.ok(gpu.source && gpu.source.length > 10, `${id} must cite a source`);
    assert.equal(gpu.maxWarpsPerSm * g.WARP_SIZE, gpu.maxThreadsPerSm, `${id} warps/threads per SM disagree`);
    assert.ok(gpu.maxSmemPerBlockKb <= gpu.l1SmemPerSmKb, `${id} block shared-mem cap exceeds the SM's total`);
    assert.ok(gpu.bandwidthGbs > 0 && gpu.l2Mb > 0);
  }
});

/* ----------------------------- Roofline ----------------------------- */

test('attainable performance is the lower of the two ceilings', () => {
  const peak = 312e12;
  const bw = 2039e9;
  // Far below the ridge: bandwidth sets the limit.
  close(g.attainableFlops(1, peak, bw), bw, 1, 'AI = 1 should be bandwidth x 1');
  // Far above it: the FLOP ceiling sets the limit.
  assert.equal(g.attainableFlops(10_000, peak, bw), peak);
  // Exactly at the ridge the two ceilings meet.
  const ridge = g.ridgePoint(peak, bw);
  close(g.attainableFlops(ridge, peak, bw), peak, peak * 1e-9);
});

test('ridge point rises across generations, because compute grew faster than bandwidth', () => {
  const ridge = (gpu) => g.ridgePoint(g.peakTensorFlops(gpu), g.peakBandwidthBytesPerSec(gpu));
  const a100 = ridge(g.GPUS.a100);
  const h100 = ridge(g.GPUS.h100);
  close(a100, 153, 1, 'A100 ridge point');
  close(h100, 295, 1, 'H100 ridge point');
  assert.ok(h100 > a100, 'the bar for being compute-bound got higher, not lower');
});

test('memory-bound classification agrees with the ridge point', () => {
  const peak = 312e12;
  const bw = 2039e9;
  assert.ok(g.isMemoryBound(10, peak, bw));
  assert.ok(!g.isMemoryBound(1000, peak, bw));
});

/* ------------------------------- GEMM ------------------------------- */

test('GEMM FLOPs and minimum traffic', () => {
  assert.equal(g.gemmFlops(2, 3, 4), 2 * 2 * 3 * 4);
  // A(2x4) + B(4x3) + C(2x3) = 8 + 12 + 6 = 26 elements, 2 bytes each.
  assert.equal(g.gemmBytes(2, 3, 4, 2), 52);
});

test('square GEMM intensity grows linearly with N', () => {
  // AI = 2N^3 / (3N^2 * bytes) = 2N / (3 * bytes)
  for (const bytes of [2, 4]) {
    for (const n of [512, 1024, 4096]) {
      close(g.gemmIntensity(n, n, n, bytes), (2 * n) / (3 * bytes), 1e-6, `N=${n}, ${bytes}B`);
    }
  }
});

test('a batch-1 decode matvec has arithmetic intensity of about 1', () => {
  const decode = g.decodeStepKernel({ dModel: 4096, bytesPerElement: 2 });
  close(decode.ai, 1, 0.01, 'every weight fetched is used exactly once');
});

/* --------------------------- Tiled matmul --------------------------- */

test('shared-memory tiling cuts global traffic by exactly the tile size', () => {
  for (const tile of [8, 16, 32, 64]) {
    const t = g.matmulTraffic(1024, tile);
    close(t.reduction, tile, 1e-9, `tile ${tile}`);
    assert.equal(t.naiveReads, 2 * 1024 ** 3);
    assert.equal(t.tiledReads, (2 * 1024 ** 3) / tile);
  }
});

test('the staged tiles must fit in shared memory', () => {
  const gpu = g.GPUS.a100;
  const fits = g.matmulTraffic(4096, 32, 4).smemBytesPerBlock; // 2 * 32*32 * 4B = 8 KB
  assert.equal(fits, 8192);
  assert.ok(fits <= gpu.maxSmemPerBlockKb * 1024, 'a 32x32 fp32 tile pair fits on an A100');

  const tooBig = g.matmulTraffic(4096, 256, 4).smemBytesPerBlock; // 512 KB
  assert.ok(tooBig > gpu.maxSmemPerBlockKb * 1024, 'a 256x256 tile pair cannot fit');
});

test('block count matches the tiling of the output matrix', () => {
  const t = g.matmulTraffic(1024, 32);
  assert.equal(t.tilesPerDim, 32);
  assert.equal(t.blocks, 32 * 32);
});

/* ---------------------------- Attention ---------------------------- */

test('fused attention intensity is n / bytesPerElement, so it grows with sequence length', () => {
  for (const n of [512, 2048, 8192]) {
    const fused = g.attentionTraffic({ seqLen: n, headDim: 128, bytesPerElement: 2, fused: true });
    close(fused.ai, n / 2, 1e-6, `n=${n}`);
  }
});

test('unfused attention intensity saturates — a longer sequence does not rescue it', () => {
  const ai = (n) => g.attentionTraffic({ seqLen: n, headDim: 128, fused: false }).ai;
  const short = ai(512);
  const long = ai(8192);
  assert.ok(long > short, 'it does improve a little');
  assert.ok(long < 64, `but it plateaus below 64 FLOP/byte (got ${long.toFixed(1)})`);

  // The plateau is d_head / 2: once n >> d_head the n^2 terms dominate traffic.
  close(ai(1e6), 128 / 2, 0.5, 'asymptote');
});

test('fusion moves attention across the A100 ridge point; not fusing never does', () => {
  const gpu = g.GPUS.a100;
  const peak = g.peakTensorFlops(gpu);
  const bw = g.peakBandwidthBytesPerSec(gpu);
  const n = 4096;
  const unfused = g.attentionTraffic({ seqLen: n, headDim: 128, fused: false });
  const fused = g.attentionTraffic({ seqLen: n, headDim: 128, fused: true });

  assert.equal(unfused.flops, fused.flops, 'fusion changes traffic, not arithmetic');
  assert.ok(g.isMemoryBound(unfused.ai, peak, bw), 'unfused sits under the roof');
  assert.ok(!g.isMemoryBound(fused.ai, peak, bw), 'fused reaches the compute ceiling');
  assert.ok(fused.bytes < unfused.bytes / 10, 'and moves an order of magnitude less data');
});

/* --------------------------- Elementwise --------------------------- */

test('elementwise kernels are memory bound by any reasonable FLOP estimate', () => {
  const gpu = g.GPUS.a100;
  const ridge = g.ridgePoint(g.peakTensorFlops(gpu), g.peakBandwidthBytesPerSec(gpu));
  // Even at a wildly generous 100 FLOP per element, GELU stays far under the ridge.
  for (const flopsPerElement of [1, 8, 100]) {
    const k = g.elementwiseKernel({ elements: 1e8, flopsPerElement, reads: 1, writes: 1 });
    assert.ok(k.ai < ridge, `${flopsPerElement} FLOP/elem gives AI ${k.ai} vs ridge ${ridge}`);
  }
});

/* ---------------------------- Occupancy ---------------------------- */

test('a modest block reaches full occupancy on an A100', () => {
  const o = g.occupancy(
    { threadsPerBlock: 256, registersPerThread: 32, sharedMemPerBlockBytes: 0 },
    g.GPUS.a100,
  );
  assert.equal(o.warpsPerBlock, 8);
  assert.equal(o.activeWarps, 64);
  assert.equal(o.occupancy, 1);
});

test('register pressure limits occupancy, and the limiter is named', () => {
  // 128 registers/thread x 32 threads = 4096 registers per warp.
  // 65536 / 4096 = 16 warps resident, out of 64 slots.
  const o = g.occupancy(
    { threadsPerBlock: 256, registersPerThread: 128, sharedMemPerBlockBytes: 0 },
    g.GPUS.a100,
  );
  assert.equal(o.regsPerWarp, 4096);
  assert.equal(o.activeWarps, 16);
  assert.equal(o.occupancy, 0.25);
  assert.deepEqual(o.limiter, ['Registers']);
});

test('registers are allocated to a warp in units of 256', () => {
  // 33 regs/thread x 32 = 1056, which rounds up to 1088+... precisely 1280.
  assert.equal(g.ceilTo(33 * 32, 256), 1280);
  const o = g.occupancy(
    { threadsPerBlock: 128, registersPerThread: 33, sharedMemPerBlockBytes: 0 },
    g.GPUS.a100,
  );
  assert.equal(o.regsPerWarp, 1280, 'the 32 wasted registers per warp are real');
});

test('shared memory can be the limiter', () => {
  // A100 has 192 KB of L1/shared per SM; 64 KB per block allows 3 blocks.
  const o = g.occupancy(
    { threadsPerBlock: 256, registersPerThread: 32, sharedMemPerBlockBytes: 64 * 1024 },
    g.GPUS.a100,
  );
  assert.equal(o.blocksPerSm, 3);
  assert.deepEqual(o.limiter, ['Shared memory']);
});

test('occupancy never exceeds 1 and respects the blocks-per-SM cap', () => {
  for (const id of g.GPU_IDS) {
    const gpu = g.GPUS[id];
    for (const threads of [32, 64, 128, 256, 512, 1024]) {
      const o = g.occupancy(
        { threadsPerBlock: threads, registersPerThread: 24, sharedMemPerBlockBytes: 0 },
        gpu,
      );
      assert.ok(o.occupancy <= 1, `${id} @ ${threads} threads gave occupancy ${o.occupancy}`);
      assert.ok(o.blocksPerSm <= gpu.maxBlocksPerSm);
      assert.ok(o.activeWarps <= gpu.maxWarpsPerSm);
    }
  }
});

test('tiny blocks waste the SM even at 100% of the block cap', () => {
  // 32-thread blocks on an A100: the 32-block cap allows only 32 of 64 warp slots.
  const o = g.occupancy(
    { threadsPerBlock: 32, registersPerThread: 24, sharedMemPerBlockBytes: 0 },
    g.GPUS.a100,
  );
  assert.equal(o.blocksPerSm, 32);
  assert.equal(o.activeWarps, 32);
  assert.equal(o.occupancy, 0.5);
  assert.deepEqual(o.limiter, ['Blocks/SM cap']);
});

/* --------------------------- Coalescing --------------------------- */

test('a contiguous warp load touches the minimum four sectors', () => {
  const c = g.coalescedTransactions(g.stridedWarpAddresses(1)); // 32 lanes x 4 B = 128 B
  assert.equal(c.sectors, 4);
  assert.equal(c.bytesRequested, 128);
  assert.equal(c.bytesFetched, 128);
  assert.equal(c.efficiency, 1);
});

test('strided access multiplies the traffic, and saturates at one sector per lane', () => {
  assert.equal(g.coalescedTransactions(g.stridedWarpAddresses(2)).sectors, 8);
  assert.equal(g.coalescedTransactions(g.stridedWarpAddresses(4)).sectors, 16);
  assert.equal(g.coalescedTransactions(g.stridedWarpAddresses(8)).sectors, 32);
  // Past stride 8 every lane already owns a sector — it cannot get worse.
  assert.equal(g.coalescedTransactions(g.stridedWarpAddresses(64)).sectors, 32);

  const worst = g.coalescedTransactions(g.stridedWarpAddresses(8));
  assert.equal(worst.wasteFactor, 8, 'fetching 1 KB to use 128 B');
});

test('all 32 lanes hitting one address is a broadcast, not 32 transactions', () => {
  const c = g.coalescedTransactions(new Array(32).fill(0));
  assert.equal(c.sectors, 1);
});

/* --------------------------- Divergence --------------------------- */

test('a uniform branch costs only the path it takes', () => {
  const allTrue = g.divergenceCost(new Array(32).fill(true), 10, 40);
  assert.equal(allTrue.divergent, false);
  assert.equal(allTrue.pathsExecuted, 1);
  assert.equal(allTrue.cycles, 10);
  assert.equal(allTrue.laneEfficiency, 1);

  const allFalse = g.divergenceCost(new Array(32).fill(false), 10, 40);
  assert.equal(allFalse.cycles, 40);
  assert.equal(allFalse.laneEfficiency, 1);
});

test('a divergent branch pays for both sides', () => {
  const lanes = Array.from({ length: 32 }, (_, i) => i < 16);
  const d = g.divergenceCost(lanes, 10, 40);
  assert.equal(d.divergent, true);
  assert.equal(d.pathsExecuted, 2);
  assert.equal(d.cycles, 50, 'sum of the paths, not the max');
  assert.equal(d.idealCycles, 40);
  // 16*10 + 16*40 = 800 useful lane-cycles out of 32*50 = 1600.
  assert.equal(d.laneEfficiency, 0.5);
});

test('one disagreeing lane is enough to serialize the whole warp', () => {
  const lanes = new Array(32).fill(false);
  lanes[7] = true;
  const d = g.divergenceCost(lanes, 100, 10);
  assert.equal(d.divergent, true);
  assert.equal(d.cycles, 110, 'the other 31 lanes wait through the expensive path');
});

/* --------------------------- Formatting --------------------------- */

test('formatters produce the units a reader expects', () => {
  assert.equal(g.formatBytes(1024), '1 KB');
  assert.equal(g.formatBytes(8192), '8 KB');
  assert.equal(g.formatBytes(1536), '1.5 KB');
  assert.equal(g.formatFlops(2e12), '2 TFLOP');
  assert.equal(g.formatCount(1.5e9), '1.5B');
});

/* ------------------ The layer catalog holds together ------------------ */

test('the transformer layer catalog computes intensities for every kernel', () => {
  const kernels = g.transformerLayerKernels({ tokens: 4096, dModel: 4096, heads: 32, seqLen: 4096 });
  assert.ok(kernels.length >= 9);
  for (const k of kernels) {
    assert.ok(Number.isFinite(k.ai) && k.ai > 0, `${k.name} has a bad AI`);
    assert.ok(k.flops > 0 && k.bytes > 0, `${k.name} has bad counts`);
    assert.equal(k.ai, k.flops / k.bytes, `${k.name}: AI must equal FLOPs/bytes`);
  }
});

test('the catalog splits into memory-bound and compute-bound as taught', () => {
  const gpu = g.GPUS.a100;
  const peak = g.peakTensorFlops(gpu);
  const bw = g.peakBandwidthBytesPerSec(gpu);
  const kernels = g.transformerLayerKernels();
  const by = (name) => kernels.find((k) => k.name === name);

  // The GEMMs are comfortably compute-bound at these shapes.
  assert.ok(!g.isMemoryBound(by('FFN up').ai, peak, bw));
  assert.ok(!g.isMemoryBound(by('QKV projection').ai, peak, bw));

  // The elementwise ops never are.
  for (const name of ['GELU', 'LayerNorm', 'Residual add']) {
    assert.ok(g.isMemoryBound(by(name).ai, peak, bw), `${name} should be memory bound`);
  }
});