Foundry/cuda_kernels/tests/bench_ssm_recurrence_seq.cu

465 lines
21 KiB
Plaintext

/* SPDX-License-Identifier: AGPL-3.0-or-later */
/* Copyright (C) 2026 David (Corey) Herrera */
/**
* @file bench_ssm_recurrence_seq.cu
* @brief Bit-exactness + timing harness for the prefill Gated-DeltaNet
* recurrence: fc_ssm_recurrence_seq (shipped baseline) vs
* fc_ssm_recurrence_seq_warp (warp-shuffle rewrite) vs
* fc_ssm_recurrence_seq_unroll128 (the cheap fallback: same barrier-chained
* sweep, literal trip count so ptxas unrolls it). Includes
* ssm_gated_delta.cu VERBATIM so every kernel is a control symbol compiled
* from the same source in the same module.
*
* Hypothesis under test: fc_ssm_recurrence_seq is ISSUE-BOUND ON OVERHEAD, not
* bandwidth-, compute-, launch- or occupancy-bound. From its committed PTX,
* per (block, position) it issues ~740 warp-instruction slots of which 66 % is
* the two shared-memory tree reductions (45 pp of that being setp/bra/shr/mov
* loop control, because the 7-step loop's bound is `state_size`, a RUNTIME
* kernel parameter, so ptxas cannot unroll it), 10 % is bar.sync, 10 % is a
* 16384x-redundant expf on a block-uniform value, and only ~5.5 % does useful
* floating-point work. Occupancy is already 100 % (512 B shared, 128-thread
* blocks, 16 blocks/SMX) and the grid is 4096 blocks, so barrier LATENCY is
* already hidden -- what is not hidden is the instruction COUNT.
*
* fc_ssm_recurrence_seq_warp cuts that to ~72 slots per (column, position) and
* removes __syncthreads from the position loop entirely. If the hypothesis is
* right, the measured ratio tracks the instruction ratio far better than it
* tracks any byte count (this project has repeatedly seen byte reduction NOT
* convert while instruction reduction did).
*
* WHAT WOULD REFUTE IT: a ratio near 1.0x despite the 10x instruction cut. That
* would mean the kernel is bound by something neither variant changes -- global
* load latency on the k/q streams, or the L1/read-only cache request rate (both
* variants issue the same 1 KB of q+k requests per column-position). In that
* case report the refutation; do not tune the warp kernel to chase it.
*
* NEGATIVE CONTROLS (this harness must be able to FAIL before green means
* anything -- see the repo's "harnesses here measure nothing" lesson):
* 1. `--self` compares the baseline against a second run of the BASELINE.
* It must print EXACT for y and s, and a ratio inside the noise floor.
* If --self ever prints DIFFER, the comparison is broken (uninitialised
* buffer, wrong length, aliasing) and every other verdict is void.
* 2. fc_ssm_recurrence_seq_warp_poison is the warp kernel with ONE deliberate
* defect: the intra-lane reduction accumulates sequentially
* (((p0+p1)+p2)+p3) instead of reproducing the baseline's tree
* ((p0+p2)+(p1+p3)). It is mathematically the same sum and numerically
* within 1e-7, so a tolerance-based check passes it. This harness MUST
* report it as DIFFER. If it reports EXACT, the bit-comparison is not
* actually comparing bits and the bit-exactness claim for the real kernel
* is unproven.
*
* Run on a K80 host (nvcc 11.4, gcc-11 host, sm_37), ONE die, no sibling load --
* two busy dies on one K80 board drop to a 398 MHz power cap:
* nvcc -arch=sm_37 -ccbin gcc-11 -Wno-deprecated-gpu-targets \
* -isystem ../nvcc_shim -O3 -Xptxas -v \
* -o bench_ssm_recurrence_seq bench_ssm_recurrence_seq.cu
* CUDA_VISIBLE_DEVICES=<free die> ./bench_ssm_recurrence_seq --self
* CUDA_VISIBLE_DEVICES=<free die> ./bench_ssm_recurrence_seq
* CUDA_VISIBLE_DEVICES=<free die> ./bench_ssm_recurrence_seq --poison
* CUDA_VISIBLE_DEVICES=<free die> ./bench_ssm_recurrence_seq --unroll
* CUDA_VISIBLE_DEVICES=<free die> ./bench_ssm_recurrence_seq --snaps
* --snaps additionally compares the per-position state SNAPSHOTS the shipped
* launcher requests in speculative-decoding mode (spec ? d->ssm_snaps : NULL).
* Without it that output of the warp kernel is never written, so the default
* arms prove nothing about it; it costs 256 MiB per side, hence the flag.
* Read the `-Xptxas -v` banner for BOTH kernels and record the STACK FRAME, not
* the spill counters: on this project an array indexed by a runtime-bounded
* loop is never promoted, and ptxas prints "0 bytes spill" right beside a live
* stack frame. The warp kernel must show 0 bytes stack frame and <= 64
* registers -- above 64 the register file caps residency below 16 blocks/SMX
* and the projected multiplier does not hold.
*
* @status Microbenchmark only. Nothing here is reachable from Pantheon; the
* shipped launcher keeps fc_ssm_recurrence_seq unless QWEN35_SSM_WARP=1.
* RUN 2026-08-10, die 6 at 875 MHz, 0.10 % noise floor: --self EXACT,
* --poison DIFFER at 1.1e-08, warp 8.167x bit-exact, unroll128 1.189x
* bit-exact, 32 registers / 0 B stack frame for all three.
* @issues
* - The warp arm requires state_size == 128 (FC_SSM_WARP_S) and stages
* 2*npos floats of shared memory, so npos must stay small (128 is the
* shipped prefill chunk).
* - Timing here is a per-kernel ratio on synthetic operands. It does NOT
* predict whole-prefill tok/s: use it to accept or reject the binding
* hypothesis, then re-measure end to end with QWEN35_PROF=1
* QWEN35_PROF_SUB=1 and read the `recur_seq` span.
* @todo None
*/
#include "../ssm_gated_delta.cu"
#include <math.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
/* Shipped prefill shape for qwen35moe / nex-n2-mini / saber-orch. */
#define BENCH_STATE_SIZE 128
#define BENCH_N_V_HEADS 32
#define BENCH_NPOS 128
#define BENCH_WARP_COLS 4
/*
* NEGATIVE CONTROL. fc_ssm_recurrence_seq_warp with exactly one defect: the
* intra-lane reduction accumulates left-to-right instead of reproducing the
* baseline's binary tree. Same value to ~1e-7, different bits. Any check that
* calls this bit-exact is not comparing bits.
*/
__device__ __forceinline__ float fc_ssm_warp_seq_dot_poison(const float *r,
const float *x) {
float acc = __fmul_rn(r[0], x[0]);
#pragma unroll
for (int e = 1; e < FC_SSM_WARP_ELEMS; e++)
acc = __fadd_rn(acc, __fmul_rn(r[e], x[e]));
#pragma unroll
for (int st = 16; st > 0; st >>= 1)
acc = __fadd_rn(acc, __shfl_xor(acc, st));
return acc;
}
extern "C" __global__ void fc_ssm_recurrence_seq_warp_poison(
float *s, const float *q, const float *k, const float *v,
const float *g, const float *beta, float *y, int state_size,
int n_v_heads, int n, float scale, float *snaps) {
extern __shared__ float sh[];
if (state_size != FC_SSM_WARP_S)
return;
const int S = FC_SSM_WARP_S;
const int lane = (int)(threadIdx.x & 31u);
const int col = (int)(threadIdx.x >> 5);
const int j = blockIdx.x * FC_SSM_WARP_COLS + col;
const int h = blockIdx.y;
float *sh_decay = sh;
float *sh_beta = sh + n;
for (int t = (int)threadIdx.x; t < n; t += FC_SSM_WARP_COLS * 32) {
sh_decay[t] = expf(g[(size_t)t * n_v_heads + h]);
sh_beta[t] = beta[(size_t)t * n_v_heads + h];
}
__syncthreads();
if (j >= S)
return;
const size_t row_off = ((size_t)h * S + j) * S;
const size_t pos_stride = (size_t)n_v_heads * S;
float *row = s + row_off;
float r[FC_SSM_WARP_ELEMS];
#pragma unroll
for (int e = 0; e < FC_SSM_WARP_ELEMS; e++)
r[e] = row[lane + e * 32];
const float *kh = k + (size_t)h * S + lane;
const float *qh = q + (size_t)h * S + lane;
const float *vj = v + (size_t)h * S + j;
float *yj = y + (size_t)h * S + j;
(void)snaps;
for (int t = 0; t < n; t++) {
const size_t toff = (size_t)t * pos_stride;
const float decay = sh_decay[t];
const float beta_h = sh_beta[t];
float kv[FC_SSM_WARP_ELEMS], qv[FC_SSM_WARP_ELEMS];
#pragma unroll
for (int e = 0; e < FC_SSM_WARP_ELEMS; e++) {
kv[e] = kh[toff + e * 32];
r[e] = __fmul_rn(r[e], decay);
}
const float pred = fc_ssm_warp_seq_dot_poison(r, kv);
const float delta = __fmul_rn(__fsub_rn(vj[toff], pred), beta_h);
#pragma unroll
for (int e = 0; e < FC_SSM_WARP_ELEMS; e++) {
r[e] = __fmaf_rn(kv[e], delta, r[e]);
qv[e] = qh[toff + e * 32];
}
const float out = fc_ssm_warp_seq_dot_poison(r, qv);
if (lane == 0)
yj[toff] = __fmul_rn(out, scale);
}
#pragma unroll
for (int e = 0; e < FC_SSM_WARP_ELEMS; e++)
row[lane + e * 32] = r[e];
}
/* ---------------------------------------------------------------- host ---- */
#define CHECK(call) \
do { \
cudaError_t rc_ = (call); \
if (rc_ != cudaSuccess) { \
fprintf(stderr, "%s:%d %s -> %s\n", __FILE__, __LINE__, #call, \
cudaGetErrorString(rc_)); \
return 1; \
} \
} while (0)
/* Deterministic operands: both arms must see byte-identical inputs or the
* bit-comparison is meaningless. */
static void fill_deterministic(float *p, size_t n, unsigned seed, float lo,
float hi) {
unsigned st = seed | 1u;
for (size_t i = 0; i < n; i++) {
st = st * 1664525u + 1013904223u;
float u = (float)(st >> 8) / (float)(1u << 24);
p[i] = lo + u * (hi - lo);
}
}
/* L2-normalise each [n_v_heads][S] row of a [npos][n_v_heads*S] stream, as the
* qk-norm kernel does upstream, so the recurrence sees realistic magnitudes. */
static void l2_norm_rows(float *p, int npos, int n_v_heads, int S) {
for (int t = 0; t < npos; t++) {
for (int h = 0; h < n_v_heads; h++) {
float *row = p + ((size_t)t * n_v_heads + h) * S;
float ss = 0.0f;
for (int i = 0; i < S; i++)
ss += row[i] * row[i];
float inv = 1.0f / fmaxf(sqrtf(ss), 1e-6f);
for (int i = 0; i < S; i++)
row[i] *= inv;
}
}
}
enum bench_arm { ARM_BASE = 0, ARM_WARP, ARM_POISON, ARM_UNROLL };
int main(int argc, char **argv) {
int self_mode = 0, poison_mode = 0, unroll_mode = 0, snaps_mode = 0, reps = 30;
for (int i = 1; i < argc; i++) {
if (!strcmp(argv[i], "--self"))
self_mode = 1;
else if (!strcmp(argv[i], "--poison"))
poison_mode = 1;
else if (!strcmp(argv[i], "--unroll"))
unroll_mode = 1;
else if (!strcmp(argv[i], "--snaps"))
snaps_mode = 1;
else if (!strcmp(argv[i], "--reps") && i + 1 < argc)
reps = atoi(argv[++i]);
}
/* Exactly two timed arms per run: the baseline, then whichever variant the
* mode selects. Two arms means the comparison buffers cannot be cross-wired
* between modes, which is how a negative control gets quietly disarmed. */
const enum bench_arm cmp_arm = self_mode ? ARM_BASE
: poison_mode ? ARM_POISON
: unroll_mode ? ARM_UNROLL
: ARM_WARP;
const int S = BENCH_STATE_SIZE, H = BENCH_N_V_HEADS, N = BENCH_NPOS;
const float scale = 1.0f / sqrtf((float)S);
const size_t n_state = (size_t)H * S * S;
const size_t n_act = (size_t)N * H * S;
const size_t n_head = (size_t)N * H;
float *h_s0 = (float *)malloc(n_state * sizeof(float));
float *h_q = (float *)malloc(n_act * sizeof(float));
float *h_k = (float *)malloc(n_act * sizeof(float));
float *h_v = (float *)malloc(n_act * sizeof(float));
float *h_g = (float *)malloc(n_head * sizeof(float));
float *h_b = (float *)malloc(n_head * sizeof(float));
float *h_ya = (float *)malloc(n_act * sizeof(float));
float *h_yb = (float *)malloc(n_act * sizeof(float));
float *h_sa = (float *)malloc(n_state * sizeof(float));
float *h_sb = (float *)malloc(n_state * sizeof(float));
if (!h_s0 || !h_q || !h_k || !h_v || !h_g || !h_b || !h_ya || !h_yb ||
!h_sa || !h_sb) {
fprintf(stderr, "host alloc failed\n");
return 1;
}
fill_deterministic(h_s0, n_state, 11u, -0.5f, 0.5f);
fill_deterministic(h_q, n_act, 22u, -1.0f, 1.0f);
fill_deterministic(h_k, n_act, 33u, -1.0f, 1.0f);
fill_deterministic(h_v, n_act, 44u, -1.0f, 1.0f);
l2_norm_rows(h_q, N, H, S);
l2_norm_rows(h_k, N, H, S);
/* g is pre-exp and negative (decay in (0,1]); beta is post-sigmoid. */
fill_deterministic(h_g, n_head, 55u, -0.30f, -0.001f);
fill_deterministic(h_b, n_head, 66u, 0.05f, 0.95f);
/* --snaps exercises the speculative-decoding output the shipped launcher
* passes when spec mode is on. It is 256 MiB per side at this shape and is
* NOT part of a timed arm's workload, so it is allocated only on request. */
const size_t n_snaps = (size_t)N * n_state;
float *h_sna = NULL, *h_snb = NULL, *d_snaps = NULL;
if (snaps_mode) {
h_sna = (float *)malloc(n_snaps * sizeof(float));
h_snb = (float *)malloc(n_snaps * sizeof(float));
if (!h_sna || !h_snb) {
fprintf(stderr, "host alloc failed (snaps)\n");
return 1;
}
}
float *d_s, *d_q, *d_k, *d_v, *d_g, *d_b, *d_y;
CHECK(cudaMalloc(&d_s, n_state * sizeof(float)));
CHECK(cudaMalloc(&d_q, n_act * sizeof(float)));
CHECK(cudaMalloc(&d_k, n_act * sizeof(float)));
CHECK(cudaMalloc(&d_v, n_act * sizeof(float)));
CHECK(cudaMalloc(&d_g, n_head * sizeof(float)));
CHECK(cudaMalloc(&d_b, n_head * sizeof(float)));
CHECK(cudaMalloc(&d_y, n_act * sizeof(float)));
if (snaps_mode)
CHECK(cudaMalloc(&d_snaps, n_snaps * sizeof(float)));
CHECK(cudaMemcpy(d_q, h_q, n_act * sizeof(float), cudaMemcpyHostToDevice));
CHECK(cudaMemcpy(d_k, h_k, n_act * sizeof(float), cudaMemcpyHostToDevice));
CHECK(cudaMemcpy(d_v, h_v, n_act * sizeof(float), cudaMemcpyHostToDevice));
CHECK(cudaMemcpy(d_g, h_g, n_head * sizeof(float), cudaMemcpyHostToDevice));
CHECK(cudaMemcpy(d_b, h_b, n_head * sizeof(float), cudaMemcpyHostToDevice));
cudaEvent_t ev0, ev1;
CHECK(cudaEventCreate(&ev0));
CHECK(cudaEventCreate(&ev1));
float ms[2] = { 0.0f, 0.0f };
for (int a = 0; a < 2; a++) {
const enum bench_arm arm = (a == 0) ? ARM_BASE : cmp_arm;
/* Only the warp variants change the launch geometry: one warp per state
* column instead of one block. The unrolled arm is the baseline's own
* geometry. */
int warp_geom = (arm == ARM_WARP || arm == ARM_POISON);
unsigned gx = warp_geom ? (unsigned)(S / BENCH_WARP_COLS) : (unsigned)S;
unsigned bx = warp_geom ? (unsigned)(BENCH_WARP_COLS * 32) : (unsigned)S;
unsigned shb = warp_geom ? (unsigned)(2 * N * sizeof(float))
: (unsigned)(S * sizeof(float));
#define LAUNCH_ARM() \
do { \
if (arm == ARM_BASE) \
fc_ssm_recurrence_seq<<<dim3(gx, H), bx, shb>>>( \
d_s, d_q, d_k, d_v, d_g, d_b, d_y, S, H, N, scale, d_snaps); \
else if (arm == ARM_WARP) \
fc_ssm_recurrence_seq_warp<<<dim3(gx, H), bx, shb>>>( \
d_s, d_q, d_k, d_v, d_g, d_b, d_y, S, H, N, scale, d_snaps); \
else if (arm == ARM_UNROLL) \
fc_ssm_recurrence_seq_unroll128<<<dim3(gx, H), bx, shb>>>( \
d_s, d_q, d_k, d_v, d_g, d_b, d_y, S, H, N, scale, d_snaps); \
else \
fc_ssm_recurrence_seq_warp_poison<<<dim3(gx, H), bx, shb>>>( \
d_s, d_q, d_k, d_v, d_g, d_b, d_y, S, H, N, scale, d_snaps); \
} while (0)
/* Pass 1 -- CORRECTNESS. Fresh state, y poisoned to 0xCD so an
* unwritten element cannot masquerade as a match, exactly one launch. */
CHECK(cudaMemcpy(d_s, h_s0, n_state * sizeof(float),
cudaMemcpyHostToDevice));
CHECK(cudaMemset(d_y, 0xCD, n_act * sizeof(float)));
if (snaps_mode)
CHECK(cudaMemset(d_snaps, 0xCD, n_snaps * sizeof(float)));
LAUNCH_ARM();
CHECK(cudaDeviceSynchronize());
CHECK(cudaGetLastError());
float *y_dst = (a == 0) ? h_ya : h_yb;
float *s_dst = (a == 0) ? h_sa : h_sb;
CHECK(cudaMemcpy(y_dst, d_y, n_act * sizeof(float),
cudaMemcpyDeviceToHost));
CHECK(cudaMemcpy(s_dst, d_s, n_state * sizeof(float),
cudaMemcpyDeviceToHost));
if (snaps_mode)
CHECK(cudaMemcpy((a == 0) ? h_sna : h_snb, d_snaps,
n_snaps * sizeof(float), cudaMemcpyDeviceToHost));
/* Pass 2 -- TIMING. Nothing but back-to-back launches inside the span:
* a per-iteration state reset would put ~4 MiB of PCIe traffic and a
* full device sync in the span and drag every ratio toward 1.00x. The
* state is reset once before the loop and then allowed to advance; with
* decay in (0,1) it stays bounded, and the instruction stream is
* identical whatever the operand values are. */
CHECK(cudaMemcpy(d_s, h_s0, n_state * sizeof(float),
cudaMemcpyHostToDevice));
for (int w = 0; w < 3; w++)
LAUNCH_ARM();
CHECK(cudaDeviceSynchronize());
CHECK(cudaEventRecord(ev0));
for (int it = 0; it < reps; it++)
LAUNCH_ARM();
CHECK(cudaEventRecord(ev1));
CHECK(cudaDeviceSynchronize());
CHECK(cudaGetLastError());
float total = 0.0f;
CHECK(cudaEventElapsedTime(&total, ev0, ev1));
ms[a] = total / (float)reps;
#undef LAUNCH_ARM
}
const char *arm_name = self_mode ? "baseline(self)"
: poison_mode ? "warp_POISON"
: unroll_mode ? "unroll128" : "warp";
size_t y_diff = 0, s_diff = 0;
double y_max = 0.0, s_max = 0.0;
for (size_t i = 0; i < n_act; i++) {
if (memcmp(&h_ya[i], &h_yb[i], sizeof(float)) != 0) {
y_diff++;
double d = fabs((double)h_ya[i] - (double)h_yb[i]);
if (d > y_max)
y_max = d;
}
}
for (size_t i = 0; i < n_state; i++) {
if (memcmp(&h_sa[i], &h_sb[i], sizeof(float)) != 0) {
s_diff++;
double d = fabs((double)h_sa[i] - (double)h_sb[i]);
if (d > s_max)
s_max = d;
}
}
printf("shape S=%d n_v_heads=%d npos=%d grid_base=(%d,%d)x%d "
"grid_warp=(%d,%d)x%d\n", S, H, N, S, H, S, S / BENCH_WARP_COLS, H,
BENCH_WARP_COLS * 32);
printf("baseline %8.4f ms/launch\n", ms[0]);
printf("%-12s %8.4f ms/launch ratio %.4fx\n", arm_name, ms[1],
ms[1] > 0.0f ? ms[0] / ms[1] : 0.0f);
printf("y bits %s (%zu/%zu differ, max |d| %.3e)\n",
y_diff ? "DIFFER" : "EXACT", y_diff, n_act, y_max);
printf("s bits %s (%zu/%zu differ, max |d| %.3e)\n",
s_diff ? "DIFFER" : "EXACT", s_diff, n_state, s_max);
size_t sn_diff = 0;
if (snaps_mode) {
double sn_max = 0.0;
for (size_t i = 0; i < n_snaps; i++) {
if (memcmp(&h_sna[i], &h_snb[i], sizeof(float)) != 0) {
sn_diff++;
double d = fabs((double)h_sna[i] - (double)h_snb[i]);
if (d > sn_max)
sn_max = d;
}
}
printf("snaps bits %s (%zu/%zu differ, max |d| %.3e)\n",
sn_diff ? "DIFFER" : "EXACT", sn_diff, n_snaps, sn_max);
}
int verdict_ok;
if (self_mode) {
verdict_ok = (y_diff == 0 && s_diff == 0);
printf("VERDICT control-1 %s -- baseline vs baseline must be "
"EXACT; ratio is the noise floor for every other run\n",
verdict_ok ? "PASS" : "FAIL (comparison is broken; all other "
"verdicts are void)");
} else if (poison_mode) {
verdict_ok = (y_diff != 0);
printf("VERDICT control-2 %s -- the poison arm MUST differ; "
"EXACT here means the bit-check is not comparing bits\n",
verdict_ok ? "PASS" : "FAIL");
} else {
verdict_ok = (y_diff == 0 && s_diff == 0 && sn_diff == 0);
printf("VERDICT %s -- the %s arm must be bit-exact in BOTH y and s; "
"run --self and --poison before trusting this\n",
verdict_ok ? "PASS" : "FAIL", arm_name);
}
free(h_s0); free(h_q); free(h_k); free(h_v); free(h_g); free(h_b);
free(h_ya); free(h_yb); free(h_sa); free(h_sb);
free(h_sna); free(h_snb); cudaFree(d_snaps);
cudaFree(d_s); cudaFree(d_q); cudaFree(d_k); cudaFree(d_v);
cudaFree(d_g); cudaFree(d_b); cudaFree(d_y);
return verdict_ok ? 0 : 2;
}