/* SPDX-License-Identifier: AGPL-3.0-or-later */ /* Copyright (C) 2026 David (Corey) Herrera */ /* * ssm_gated_delta.cu — Qwen3.5-MoE Gated-Delta-Net SSM kernels for Foundry * * Target: sm_37 (K80, Kepler) * Compile: nvcc -ptx -arch=sm_37 ssm_gated_delta.cu -o ssm_gated_delta.ptx * * The GPU halves of the single-token (decode) Gated-Delta-Net SSM block. Each * kernel mirrors, line-for-line, a static helper in the CPU reference * Pantheon/src/models/qwen35moe_ssm.c * and is validated against the golden vectors in Pantheon/tests/golden/. * * Kernels emitted (PTX symbols for ../cuda/fc_cuda_ptx.h, FC_PTX_SSM_GATED_DELTA). * ★ ADDING A KERNEL HERE IS A THREE-STEP CHANGE, and skipping either of the last * two makes the kernel INERT while the serve prints a misleading load failure: * 1. define it here, then * 2. `cd cuda_kernels && make` ON A HOST WITH nvcc 11.4 to regenerate ../cuda/fc_cuda_ptx.h * (a checked-in artifact — the kernel does not exist until it is in there), and * 3. add its name to ssm_names[] in fc_cuda.c AND bump the count argument to * fc_load_ptx_module. The loader assigns names in order and stops at the * first one the module lacks, so a name with no symbol behind it makes * fc_load_ptx_module return nonzero and print "SSM gated-delta PTX load * failed (kernels unavailable)" even though every earlier kernel loaded fine. * The oldest three, for orientation (the rest are documented at their definitions): * fc_ssm_recurrence — per-token delta-rule recurrence (THE centerpiece) * mirrors qwen35moe_compute_recur_heads (ssm.c:163-196) * fc_ssm_conv1d_silu — causal depthwise conv1d (k=4) + SiLU, rolls conv_state * mirrors qwen35moe_conv1d_silu (ssm.c:110-135) * fc_ssm_gated_norm — gated RMSNorm: RMSNorm(y,w) * silu(z), per V-head * mirrors qwen35moe_gated_norm (ssm.c:206-223) * * Config (verified Nex-N2-mini): n_v_heads=32, n_k_heads=16, state_size S=128, * scale=1/sqrt(128), conv_dim=8192, conv_kernel=4. Transposed state layout: * s[h*S*S + j*S + i] == S_h[i][j] (row j contiguous over i). * The q/k arriving here are ALREADY L2-normed and GQA-repeated to n_v_heads on * the host (interleaved tiling vh = r*n_k_heads + kh); this kernel never sees * the 16->32 repeat, only the post-repeat [n_v_heads*S] streams. * * Kepler notes: pure fp32 throughout; plain shared-memory tree reductions * (no __shfl / _sync portability question); expf/__expf via the standard libm * device intrinsics, matching the CPU expf bit-tolerance band (~1e-4..1e-3 rel). */ extern "C" { /* SiLU(x) = x * sigmoid(x) = x / (1 + exp(-x)). * Mirrors qwen35moe_silu (ssm.c:60-63). */ __device__ __forceinline__ float fc_ssm_silu(float x) { return x / (1.0f + expf(-x)); } /* * ============================================================ * Kernel 1: per-token delta-rule recurrence (THE centerpiece) * ============================================================ * * Mirrors qwen35moe_compute_recur_heads (ssm.c:163-196), one (head h, row j) * per CUDA block. The CPU walks all S rows of all heads serially; the GPU maps * the 4096 independent (h,j) pairs onto the grid. Within a block the S=128 * threads cooperate on the two length-S dot reductions. * * State is transposed: s[h*S*S + j*S + i] == S_h[i][j]; row j is contiguous * over i, so each block owns one unit-stride row `row[i] = s[h*S*S + j*S + i]`. * For one token (CPU lines 180-193): * (a) row[i] *= exp(g[h]) scalar decay of the whole row * (b) pred = sum_i row[i]*k[h][i] (shared-mem tree reduction) * delta = (v[h][j] - pred) * beta[h] * (c) row[i] += k[h][i] * delta rank-1 write, in place * (d) y[h][j] = scale * sum_i row[i]*q[h][i] read-AFTER-write (2nd reduce) * The decayed+updated row is written back to s in place (state carries forward). * * Inputs: * s [n_v_heads * S * S] transposed state, updated in place * q,k,v [n_v_heads * S] q/k L2-normed + GQA-repeated on host; v straight * g [n_v_heads] scalar decay (PRE-exp; we exp here, CPU line 173) * beta [n_v_heads] sigmoid'd on host (CPU applies sigmoid earlier) * y [n_v_heads * S] output, one scalar per (h,j) * state_size S (must equal blockDim.x) * scale 1/sqrt(S) * * Grid: (state_size, n_v_heads, 1) — blockIdx.x = j (row), blockIdx.y = h. * Block: state_size threads (= S = 128), one thread per i. * Shared mem: blockDim.x * sizeof(float) (reduction scratch). */ __global__ void fc_ssm_recurrence(float *s, const float *q, const float *k, const float *v, const float *g, const float *beta, float *y, int state_size, float scale) { extern __shared__ float sdata[]; const int S = state_size; const int j = blockIdx.x; /* row j of the transposed state */ const int h = blockIdx.y; /* V-head */ const int i = threadIdx.x; /* state coordinate within the row */ if (i >= S) return; float *row = s + ((size_t)h * S + j) * S; /* s[h*S*S + j*S + i] over i */ const float *kh = k + (size_t)h * S; const float *qh = q + (size_t)h * S; const float decay = expf(g[h]); const float beta_h = beta[h]; /* (a) scalar decay of the whole row (CPU lines 180-182). */ float ri = row[i] * decay; /* (b) pred = sum_i row'[i]*k[i] via shared-mem tree reduction. */ sdata[i] = ri * kh[i]; __syncthreads(); for (int stride = S / 2; stride > 0; stride >>= 1) { if (i < stride) sdata[i] += sdata[i + stride]; __syncthreads(); } float pred = sdata[0]; /* delta = (v[h][j] - pred) * beta (CPU line 186; v indexed by j). */ float delta = (v[(size_t)h * S + j] - pred) * beta_h; __syncthreads(); /* all threads consumed sdata[0] before the next reduce */ /* (c) rank-1 write: row[i] += k[i]*delta (CPU lines 188-190), in place. */ ri += kh[i] * delta; row[i] = ri; /* (d) read-after-write: y[h][j] = scale * sum_i row[i]*q[i] (CPU line 193). */ sdata[i] = ri * qh[i]; __syncthreads(); for (int stride = S / 2; stride > 0; stride >>= 1) { if (i < stride) sdata[i] += sdata[i + stride]; __syncthreads(); } if (i == 0) y[(size_t)h * S + j] = scale * sdata[0]; } /* * ============================================================ * Kernel 2: causal conv1d (k=4) + SiLU, rolls conv_state * ============================================================ * * Mirrors qwen35moe_conv1d_silu (ssm.c:110-135). One thread per channel c over * the full conv_dim stream (depthwise; channels are independent). Each channel's * sliding window is [conv_state(c,0) .. conv_state(c,k-2), qkv(c)] (length k): * out[c] = SiLU( sum_{t= conv_dim) return; const int k = conv_kernel; const int km1 = k - 1; const float *cw = weight + (size_t)c * k; float *cs = conv_state + (size_t)c * km1; /* out[c] = SiLU( history . w[0..k-2] + qkv[c]*w[k-1] ) (CPU lines 125-128). */ float acc = qkv[c] * cw[km1]; for (int t = 0; t < km1; t++) acc += cs[t] * cw[t]; out[c] = fc_ssm_silu(acc); /* Roll the per-channel state: drop oldest, append raw qkv (CPU lines 131-133). */ for (int t = 0; t < km1 - 1; t++) cs[t] = cs[t + 1]; cs[km1 - 1] = qkv[c]; } /* * ============================================================ * Kernel 3: gated RMSNorm — RMSNorm(y, w) * silu(z), per V-head * ============================================================ * * Mirrors qwen35moe_gated_norm (ssm.c:206-223), minus the trailing ssm_out * projection (that reuses fc_matvec_q3k on the host). One block per V-head; the * S=128 threads do a shared-mem sum-of-squares reduction, then each writes * gated[h*S+d] = (y[h*S+d] / sqrt(mean(y^2)+eps)) * w[d] * silu(z[h*S+d]). * The norm weight w is shared across heads (length S), exactly like the CPU * pn_vrms_norm_f32 call per head (CPU line 219) followed by the silu(z) multiply * (CPU lines 220-221). * * Inputs: * y [n_v_heads * S] recurrence output * weight [S] ssm_norm weight (shared across heads) * z [n_v_heads * S] attn_gate projection (pre-silu) * gated [n_v_heads * S] output * state_size S, eps * * Grid: (n_v_heads, 1, 1) one block per head. * Block: 256 threads (or S if smaller). Shared mem: blockDim.x * sizeof(float). */ __global__ void fc_ssm_gated_norm(const float *y, const float *weight, const float *z, float *gated, int state_size, float eps) { extern __shared__ float sdata[]; const int S = state_size; const int h = blockIdx.x; const int tid = threadIdx.x; const int threads = blockDim.x; const float *yh = y + (size_t)h * S; const float *zh = z + (size_t)h * S; float *gh = gated + (size_t)h * S; /* Phase 1: partial sum of squares of y over this head's S dims. */ float sum_sq = 0.0f; for (int d = tid; d < S; d += threads) { float v = yh[d]; sum_sq += v * v; } sdata[tid] = sum_sq; __syncthreads(); for (int stride = threads / 2; stride > 0; stride >>= 1) { if (tid < stride) sdata[tid] += sdata[tid + stride]; __syncthreads(); } /* RMSNorm divisor matches pn_vrms_norm_f32: sqrt(mean(y^2) + eps). */ float inv_rms = 1.0f / sqrtf(sdata[0] / (float)S + eps); /* Phase 2: gated[d] = RMSNorm(y)[d] * w[d] * silu(z[d]). */ for (int d = tid; d < S; d += threads) gh[d] = yh[d] * inv_rms * weight[d] * fc_ssm_silu(zh[d]); } /* ================================================================ * SSM "glue" kernels — replace the host round-trips in the qwen35moe GPU * forward (each was a D2H -> CPU -> H2D per SSM layer, ~270 sync'd PCIe * transfers/token = the tok/s floor). Small element-wise ops, pure fp32. * ================================================================ */ /* Per-element sigmoid, in place: x = 1/(1+exp(-x)). One thread per element. * Replaces the host-side beta = sigmoid(beta). */ __global__ void fc_sigmoid(float *x, int n) { int i = blockIdx.x * blockDim.x + threadIdx.x; if (i < n) x[i] = 1.0f / (1.0f + expf(-x[i])); } /* GDA decay gate: g[h] = ssm_a[h] * softplus(alpha[h] + ssm_dt[h]), per V-head. * softplus(v) = v>20 ? v : log1p(exp(v)); ssm_a = -exp(A_log) (pre-baked host * side). One thread per head. Replaces the 3x D2H + CPU + H2D decay-gate glue. */ __global__ void fc_ssm_decay_gate(const float *alpha, const float *ssm_dt, const float *ssm_a, float *g, int n) { int h = blockIdx.x * blockDim.x + threadIdx.x; if (h >= n) return; float v = alpha[h] + ssm_dt[h]; float sp = (v > 20.0f) ? v : log1pf(expf(v)); g[h] = ssm_a[h] * sp; } /* Per-K-head L2-norm of q and k over their S state dims, then GQA-repeat * interleaved to n_v_heads (vh = r*n_k_heads + kh). conv_out layout is * [q key_dim | k key_dim | v val_dim]; this handles q,k only (v is a straight * copy done by the launcher). One block per K-head, blockDim.x == S, shared mem * [2*S]. L2-norm matches gpu_l2_norm: x / max(sqrt(sum x^2), eps). */ __global__ void fc_ssm_qknorm_gqa(const float *conv_out, float *qg, float *kg, int n_k_heads, int kv_groups, int S, int key_dim, float eps) { int kh = blockIdx.x; if (kh >= n_k_heads) return; int i = threadIdx.x; /* one thread per state dim (== S) */ extern __shared__ float sh[]; /* [2*blockDim.x]: q | k sumsq partials */ const float *q_src = conv_out; const float *k_src = conv_out + key_dim; float qv = q_src[(size_t)kh * S + i]; float kv = k_src[(size_t)kh * S + i]; sh[i] = qv * qv; sh[blockDim.x + i] = kv * kv; __syncthreads(); for (int stride = blockDim.x >> 1; stride > 0; stride >>= 1) { if (i < stride) { sh[i] += sh[i + stride]; sh[blockDim.x + i] += sh[blockDim.x + i + stride]; } __syncthreads(); } float inv_q = 1.0f / fmaxf(sqrtf(sh[0]), eps); float inv_k = 1.0f / fmaxf(sqrtf(sh[blockDim.x]), eps); float qn = qv * inv_q; float kn = kv * inv_k; for (int r = 0; r < kv_groups; r++) { int vh = r * n_k_heads + kh; /* interleaved GQA tiling */ qg[(size_t)vh * S + i] = qn; kg[(size_t)vh * S + i] = kn; } } /* ================================================================ * BATCHED (prefill) variants — n positions of ONE sequence per launch. * * The two stateful kernels (conv1d, recurrence) are a SCAN over the * positions: each block loops t = 0..n-1 internally, carrying its slice of * the state forward exactly as n single-token launches would — the per- * position fp op order is identical to the single-token kernels above, so * every position's output and the final carried state are bit-exact. The * stateless glue (decay gate, qknorm+GQA) just gains a position axis. * * Batched activations are token-major: qkv/out [n][conv_dim], * q/k/v/y [n][n_v_heads*S], g/beta [n][n_v_heads]. * ================================================================ */ /* Causal conv1d + SiLU over n positions (scan). One thread per channel; * the (k-1)-deep history lives in registers across positions and is * written back once at the end. Per position the accumulation order * matches fc_ssm_conv1d_silu exactly. Requires conv_kernel <= 4. * * snaps (optional, NULL = off): per-position state snapshots for * speculative decoding — snaps[t] = the conv_state image AFTER position t * was consumed, [n][conv_dim][km1] (same channel-major layout as * conv_state), so a partial accept of M positions restores state-after-M * with one contiguous copy of snaps[M-1]. * * LOCAL-MEMORY NOTE (why the k==4 case below is hand-scalarised): * the generic body keeps the history in `float hist[3]` indexed by loops * bounded by the RUNTIME parameter km1, and stores hist[km1-1] at a runtime * index. nvcc cannot promote an array under a runtime-indexed store, so the * whole array went to local memory — the committed PTX carries * .local .align 4 .b8 __local_depot6[12]; * with ~3 ld.local + 3 st.local per position per thread plus two * unroll-remainder branch ladders, while `ptxas -v` still prints * "0 bytes spill" (read the STACK FRAME, not the spill counters). * conv_kernel is 4 for every model that reaches here and the launcher * already rejects conv_kernel > 4, so k==4 is specialised into three named * scalars: no runtime index, no local traffic. BIT-EXACT — the accumulation is * the same `mul` then three `fma`s in the same u = 0,1,2 order the generic path * already emitted (verified against the committed PTX). * * MEASURED IN THE REGENERATED PTX (2026-08-10), and read this before chasing the * frame again: `__local_depot6[12]` is STILL DECLARED, because the generic k<4 * fallback below lives in the same entry and still stores at a runtime index, so * nvcc must reserve the frame for it. What moved is the *executed* path — the * `setp.eq %r, 4` branch target carries **0 ld.local and 0 st.local**, against * **26 ld.local + 14 st.local** when k==4 had no specialisation and that was the * only path there was. Every model reaching here has conv_kernel == 4, so no * executed instruction touches local memory and the reserved frame costs nothing * but address space. Getting the frame to a literal zero needs the k==2 and k==3 * cases hand-scalarised too (or the launcher narrowed to conv_kernel == 4 and the * generic body deleted); this kernel is 0.4 % of prefill, so that was not done. */ __global__ void fc_ssm_conv1d_silu_seq(const float *qkv, const float *weight, float *conv_state, float *out, int conv_dim, int conv_kernel, int n, float *snaps) { int c = blockIdx.x * blockDim.x + threadIdx.x; if (c >= conv_dim) return; const int k = conv_kernel; const int km1 = k - 1; const float *cw = weight + (size_t)c * k; float *cs = conv_state + (size_t)c * km1; if (k == 4) { const float w0 = cw[0], w1 = cw[1], w2 = cw[2], w3 = cw[3]; float h0 = cs[0], h1 = cs[1], h2 = cs[2]; for (int t = 0; t < n; t++) { float xv = qkv[(size_t)t * conv_dim + c]; float acc = __fmul_rn(xv, w3); acc = __fmaf_rn(h0, w0, acc); acc = __fmaf_rn(h1, w1, acc); acc = __fmaf_rn(h2, w2, acc); out[(size_t)t * conv_dim + c] = fc_ssm_silu(acc); h0 = h1; h1 = h2; h2 = xv; if (snaps) { float *sn = snaps + ((size_t)t * conv_dim + c) * 3; sn[0] = h0; sn[1] = h1; sn[2] = h2; } } cs[0] = h0; cs[1] = h1; cs[2] = h2; return; } float hist[3]; /* km1 <= 3 (conv_kernel <= 4) */ for (int t = 0; t < km1; t++) hist[t] = cs[t]; for (int t = 0; t < n; t++) { float xv = qkv[(size_t)t * conv_dim + c]; float acc = xv * cw[km1]; for (int u = 0; u < km1; u++) acc += hist[u] * cw[u]; out[(size_t)t * conv_dim + c] = fc_ssm_silu(acc); for (int u = 0; u < km1 - 1; u++) hist[u] = hist[u + 1]; hist[km1 - 1] = xv; if (snaps) { float *sn = snaps + ((size_t)t * conv_dim + c) * km1; for (int u = 0; u < km1; u++) sn[u] = hist[u]; } } for (int t = 0; t < km1; t++) cs[t] = hist[t]; } /* Delta-rule recurrence over n positions (scan). Block (j, h) loops the * positions, carrying its state row in registers (ri) between them — the * row is read once, updated n times with the exact per-position op order * of fc_ssm_recurrence, and written back once. q/k/v/y advance by * n_v_heads*S per position, g/beta by n_v_heads. * * snaps (optional, NULL = off): per-position state snapshots for * speculative decoding — snaps[t] = the full recurrent state AFTER * position t's update, [n][n_v_heads*S*S] (same layout as s), so a * partial accept of M positions restores state-after-M with one * contiguous copy of snaps[M-1]. */ __global__ void fc_ssm_recurrence_seq(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 sdata[]; const int S = state_size; const int j = blockIdx.x; const int h = blockIdx.y; const int i = threadIdx.x; if (i >= S) return; const size_t row_off = ((size_t)h * S + j) * S; const size_t state_total = (size_t)n_v_heads * S * S; float *row = s + row_off; float ri = row[i]; for (int t = 0; t < n; t++) { const float *kh = k + (size_t)t * n_v_heads * S + (size_t)h * S; const float *qh = q + (size_t)t * n_v_heads * S + (size_t)h * S; const float decay = expf(g[(size_t)t * n_v_heads + h]); const float beta_h = beta[(size_t)t * n_v_heads + h]; ri *= decay; sdata[i] = ri * kh[i]; __syncthreads(); for (int stride = S / 2; stride > 0; stride >>= 1) { if (i < stride) sdata[i] += sdata[i + stride]; __syncthreads(); } float pred = sdata[0]; float delta = (v[(size_t)t * n_v_heads * S + (size_t)h * S + j] - pred) * beta_h; __syncthreads(); ri += kh[i] * delta; if (snaps) snaps[(size_t)t * state_total + row_off + i] = ri; sdata[i] = ri * qh[i]; __syncthreads(); for (int stride = S / 2; stride > 0; stride >>= 1) { if (i < stride) sdata[i] += sdata[i + stride]; __syncthreads(); } if (i == 0) y[(size_t)t * n_v_heads * S + (size_t)h * S + j] = scale * sdata[0]; __syncthreads(); } row[i] = ri; } /* UNROLLED-REDUCTION prefill recurrence — the cheap half of the warp rewrite. * * WHY. fc_ssm_recurrence_seq's two tree reductions are written * `for (stride = S/2; stride > 0; stride >>= 1)` where S == state_size is a * RUNTIME kernel parameter, so ptxas cannot unroll the 7-step tree and every * step pays a predicate, a branch, a shift and a latch. From the committed PTX * that loop control alone is 336 of the kernel's ~740 warp-issue slots per * (block, position) — 45 %. Specialising S to 128 makes the bound a literal and * lets ptxas unroll both trees, deleting the loop control but keeping the * barrier chain and the shared-memory traffic: ~1.83x fewer issue slots, against * ~10x for fc_ssm_recurrence_seq_warp below. * * It exists as a separately priceable arm: it is ~15 lines and cannot regress * occupancy (identical shared and register footprint), so if the warp rewrite * fails its ptxas gate this is the fallback that still collects the loop-control * half of the win. * * BIT-EXACT, and trivially so: this is fc_ssm_recurrence_seq verbatim with the * loop bound written as a literal. Every floating-point operation is the same * expression in the same order with the same operands, and every intermediate * still round-trips through shared memory (so ptxas cannot contract a product * into an fma here either). Do NOT "tidy" the plain `*` and `+=` into * intrinsics: matching the baseline's codegen is the point. * * Guarded on state_size == FC_SSM_UNROLL_S; the guard is block-uniform so it can * never split a __syncthreads. The launcher keeps the baseline for anything else. * * Grid/block/shared: identical to fc_ssm_recurrence_seq — (S, n_v_heads), * S threads, S floats of shared. */ #define FC_SSM_UNROLL_S 128 /* state_size this variant is specialised for */ __global__ void fc_ssm_recurrence_seq_unroll128(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 sdata[]; /* Block-uniform: every thread returns or none does. */ if (state_size != FC_SSM_UNROLL_S) return; const int S = FC_SSM_UNROLL_S; const int j = blockIdx.x; const int h = blockIdx.y; const int i = threadIdx.x; if (i >= S) return; const size_t row_off = ((size_t)h * S + j) * S; const size_t state_total = (size_t)n_v_heads * S * S; float *row = s + row_off; float ri = row[i]; for (int t = 0; t < n; t++) { const float *kh = k + (size_t)t * n_v_heads * S + (size_t)h * S; const float *qh = q + (size_t)t * n_v_heads * S + (size_t)h * S; const float decay = expf(g[(size_t)t * n_v_heads + h]); const float beta_h = beta[(size_t)t * n_v_heads + h]; ri *= decay; sdata[i] = ri * kh[i]; __syncthreads(); #pragma unroll for (int stride = FC_SSM_UNROLL_S / 2; stride > 0; stride >>= 1) { if (i < stride) sdata[i] += sdata[i + stride]; __syncthreads(); } float pred = sdata[0]; float delta = (v[(size_t)t * n_v_heads * S + (size_t)h * S + j] - pred) * beta_h; __syncthreads(); ri += kh[i] * delta; if (snaps) snaps[(size_t)t * state_total + row_off + i] = ri; sdata[i] = ri * qh[i]; __syncthreads(); #pragma unroll for (int stride = FC_SSM_UNROLL_S / 2; stride > 0; stride >>= 1) { if (i < stride) sdata[i] += sdata[i + stride]; __syncthreads(); } if (i == 0) y[(size_t)t * n_v_heads * S + (size_t)h * S + j] = scale * sdata[0]; __syncthreads(); } row[i] = ri; } /* ================================================================ * WARP-REDUCTION prefill recurrence — the tree-preserving rewrite. * * WHY. In fc_ssm_recurrence_seq above, one 128-thread block advances ONE * state column (h, j) and spends its whole life in two shared-memory tree * reductions. From the committed PTX for that kernel, per (block, position): * * 740 warp-instruction issue slots, of which * 488 (66 %) the two tree reductions, and inside those * 336 (45 %) setp/bra/shr/mov loop control — the 7-step loop is NOT * unrolled because its bound is `state_size`, a RUNTIME * kernel parameter, so every step pays a predicate, a * branch, a shift and a latch * 72 ( 10 %) bar.sync (18 barriers per position, not 14) * 48 ( 6 %) the actual ld.shared/add.f32/st.shared * 76 (10 %) expf(g[t][h]) — 19 PTX ops incl. ex2.approx.ftz, on a * value UNIFORM across the entire block, recomputed by all * 128 threads at all 128 positions (16384x redundant) * and only ~41 slots (5.5 %) doing useful floating-point work. * Occupancy is already 100 % (512 B shared, ~40 regs, 16 blocks/SMX), the * grid is 4096 blocks = 315 per SMX, and the kernel still sustains only * IPC ~2/SMX. It is not bandwidth-, compute-, launch- or occupancy-bound: * it is issue-bound on overhead, and 3 of its 4 warps idle at a barrier * through 5 of the 7 tree steps. * * WHAT. One WARP owns one column (h, j); lane L carries state elements * i = L, L+32, L+64, L+96 in registers. The length-128 dot then costs 4 * products, 3 intra-lane adds and a 5-step __shfl_xor butterfly — no shared * memory, and __syncthreads() disappears from the position loop entirely * (one barrier per BLOCK, for the decay/beta staging). ~72 issue slots per * (column, position) against 740: a 10x instruction cut with 49 % of the * slots now doing useful FP work. * * BIT-EXACT, and this is the whole point of the structure: * - fc_ssm_warp_tree_dot reproduces the baseline's reduction TREE, not just * its value. The baseline pairs index i with i+64, then i with i+32, then * runs strides 16..1. In lane-local terms 64 and 32 are e-strides 2 and 1 * over the 4 elements a lane holds, so p[0]+=p[2]; p[1]+=p[3]; p[0]+=p[1] * lands the identical grouping, and strides 16..1 are exactly a 32-lane * butterfly. __shfl_xor gives lane L the pair (v_L, v_{L^st}); for L < st * that is the baseline's (v_L, v_{L+st}) and for L >= st it is the same * two values with the operands swapped — fp add is commutative bit-for-bit, * so every lane ends holding the baseline's sdata[0] exactly, which also * removes the broadcast. * - every arithmetic op uses __fmul_rn/__fadd_rn/__fsub_rn/__fmaf_rn. This * is NOT decoration: the baseline's per-element product went through * shared memory, so nvcc could not contract it, while here everything is * in registers and nvcc WOULD fuse p[0] + p[2] into fma(r2, k2, p0) and * silently change the result. The _rn intrinsics are contraction barriers. * - expf is hoisted, not approximated: the same expf() on the same fp32 * input, stored fp32 in shared and reloaded, is the same bits. * * Specialised to state_size == 128 (every model that reaches this kernel); * the launcher keeps the baseline for anything else, and the guard below is * block-uniform so it can never split a __syncthreads. * * Grid: (state_size / FC_SSM_WARP_COLS, n_v_heads) = (32, 32) = 1024 blocks * Block: FC_SSM_WARP_COLS * 32 = 128 threads (4 warps = 4 columns) * Shared: 2 * n floats — decay[n] | beta[n], staged once per block. * ================================================================ */ #define FC_SSM_WARP_S 128 /* state_size this variant is specialised for */ #define FC_SSM_WARP_ELEMS 4 /* FC_SSM_WARP_S / 32 state elements per lane */ #define FC_SSM_WARP_COLS 4 /* warps per block == state columns per block */ #if FC_SSM_WARP_ELEMS != 4 #error "fc_ssm_warp_tree_dot's pairing is derived for S=128 (4 elems/lane)" #endif /* Length-FC_SSM_WARP_S dot of the lane-distributed vectors r and x, reducing * in the exact binary tree fc_ssm_recurrence_seq's shared-memory sweep uses. * Returns the identical fp32 bits in all 32 lanes. * * p[] is indexed only by loops with COMPILE-TIME trip counts, so it is * register-promoted. Verify that, do not assume it: `ptxas -v` must report a * ZERO stack frame for this kernel. On this project a small array under a * runtime-bounded loop lands in local memory while the spill counters still * read "0 bytes spill" (see fc_ssm_conv1d_silu_seq's __local_depot above). */ __device__ __forceinline__ float fc_ssm_warp_tree_dot(const float *r, const float *x) { /* Level 0: the products the baseline writes into sdata[i]. */ float p[FC_SSM_WARP_ELEMS]; #pragma unroll for (int e = 0; e < FC_SSM_WARP_ELEMS; e++) p[e] = __fmul_rn(r[e], x[e]); /* Baseline strides S/2 .. 32 — entirely inside one lane (e-strides). */ #pragma unroll for (int st = FC_SSM_WARP_ELEMS / 2; st > 0; st >>= 1) #pragma unroll for (int e = 0; e < st; e++) p[e] = __fadd_rn(p[e], p[e + st]); /* Baseline strides 16 .. 1 over the 32 survivors — a warp butterfly. */ float acc = p[0]; #pragma unroll for (int st = 16; st > 0; st >>= 1) acc = __fadd_rn(acc, __shfl_xor(acc, st)); return acc; } __global__ void fc_ssm_recurrence_seq_warp(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[]; /* [2*n]: decay[n] | beta[n] */ /* Block-uniform: every thread returns or none does. */ 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; /* Stage the block-uniform decay/beta for all n positions once. */ 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(); /* the only barrier in this kernel */ if (j >= S) return; const size_t row_off = ((size_t)h * S + j) * S; const size_t state_total = (size_t)n_v_heads * S * 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; 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]; #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_tree_dot(r, kv); const float delta = __fmul_rn(__fsub_rn(vj[toff], pred), beta_h); float qv[FC_SSM_WARP_ELEMS]; #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]; } if (snaps) { float *sn = snaps + (size_t)t * state_total + row_off; #pragma unroll for (int e = 0; e < FC_SSM_WARP_ELEMS; e++) sn[lane + e * 32] = r[e]; } const float out = fc_ssm_warp_tree_dot(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]; } /* ================================================================ * MULTI-USER DECODE (slot) variants — K independent sequences, ONE token * each, per launch. Unlike the _seq scans above there is NO cross-position * dependency: slot u owns its OWN state image at offset u and advances it * exactly one step. Per slot the fp op order is identical to the * single-token kernels (fc_ssm_conv1d_silu / fc_ssm_recurrence), so each * user's output and carried state are bit-exact vs decoding alone. * * Slot-major activations match the batch layout: qkv/out [K][conv_dim], * q/k/v/y [K][n_v_heads*S], g/beta [K][n_v_heads]. States are slot-major: * conv_state [K][conv_dim*(k-1)], s [K][n_v_heads*S*S]. * ================================================================ */ /* Causal conv1d + SiLU, one step per slot. Grid (ceil(conv_dim/block), K); * per (slot, channel) the math is fc_ssm_conv1d_silu verbatim on that * slot's state slice. */ __global__ void fc_ssm_conv1d_silu_slots(const float *qkv, const float *weight, float *conv_state, float *out, int conv_dim, int conv_kernel) { int c = blockIdx.x * blockDim.x + threadIdx.x; int u = blockIdx.y; /* slot (user) */ if (c >= conv_dim) return; const int k = conv_kernel; const int km1 = k - 1; const float *cw = weight + (size_t)c * k; float *cs = conv_state + ((size_t)u * conv_dim + c) * km1; const float *x = qkv + (size_t)u * conv_dim; float *o = out + (size_t)u * conv_dim; float acc = x[c] * cw[km1]; for (int t = 0; t < km1; t++) acc += cs[t] * cw[t]; o[c] = fc_ssm_silu(acc); for (int t = 0; t < km1 - 1; t++) cs[t] = cs[t + 1]; cs[km1 - 1] = x[c]; } /* Delta-rule recurrence, one step per slot. Grid (S, n_v_heads, K); block * (j, h, u) updates row j of head h of SLOT u's state — the body is * fc_ssm_recurrence verbatim at that slot's offsets. */ __global__ void fc_ssm_recurrence_slots(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, float scale) { extern __shared__ float sdata[]; const int S = state_size; const int j = blockIdx.x; /* row j of the transposed state */ const int h = blockIdx.y; /* V-head */ const int u = blockIdx.z; /* slot (user) */ const int i = threadIdx.x; /* state coordinate within the row */ if (i >= S) return; float *row = s + (size_t)u * n_v_heads * S * S + ((size_t)h * S + j) * S; const float *kh = k + (size_t)u * n_v_heads * S + (size_t)h * S; const float *qh = q + (size_t)u * n_v_heads * S + (size_t)h * S; const float decay = expf(g[(size_t)u * n_v_heads + h]); const float beta_h = beta[(size_t)u * n_v_heads + h]; float ri = row[i] * decay; sdata[i] = ri * kh[i]; __syncthreads(); for (int stride = S / 2; stride > 0; stride >>= 1) { if (i < stride) sdata[i] += sdata[i + stride]; __syncthreads(); } float pred = sdata[0]; float delta = (v[(size_t)u * n_v_heads * S + (size_t)h * S + j] - pred) * beta_h; __syncthreads(); ri += kh[i] * delta; row[i] = ri; sdata[i] = ri * qh[i]; __syncthreads(); for (int stride = S / 2; stride > 0; stride >>= 1) { if (i < stride) sdata[i] += sdata[i + stride]; __syncthreads(); } if (i == 0) y[(size_t)u * n_v_heads * S + (size_t)h * S + j] = scale * sdata[0]; } /* ================================================================ * T4: 3-pass -> 2-pass state-access factorization (fable2 member T4). * * The baseline fc_ssm_recurrence_slots touches the recurrent state twice in * global memory (one read of row[i], one write), but it runs TWO shared-memory * tree reductions per (h,j) block separated by a read-after-write barrier: * R1: pred = sum_i (decay*row_prev[i]) * k[i] -> delta * R2: y[j] = scale * sum_i row_new[i] * q[i] (row_new needs delta) * R2 is data-dependent on R1 (it reads the just-written row), so on Kepler the * two length-S reductions serialize behind ~2*log2(S) __syncthreads barriers * plus a second LDS sweep — the dominant per-block cost. * * The T4 algebraic identity S_t^T q = g*(S_{t-1}^T q) + (q^T k)*delta lets the * OUTPUT ride the SAME read pass as retrieval: with row_dec[i] = decay*row_prev[i], * y[j] = scale * sum_i (row_dec[i] + k[i]*delta) * q[i] * = scale * ( sum_i row_dec[i]*q[i] + delta * sum_i k[i]*q[i] ) * = scale * ( oq[j] + delta * qk ) * where qk = sum_i k[i]*q[i] is a scalar shared across all rows j of the head. * All three partials (pred = sum ri*k, oq = sum ri*q, qk = sum k*q) are reduced * together in ONE fused pass ([3*S] shared), removing the second reduction and * its read-after-write barrier. State traffic is unchanged (still one read, one * write); the win is the eliminated LDS sweep + barriers. * * LOSSY-tiny (reassociation class): the baseline forms y as one fused dot over * row_new[i]; this forms it as two separate sums (oq + delta*qk). The result * differs only by FP associativity of sum((a+b)*q) vs sum(a*q)+sum(b*q) — * same magnitude as any tile/reduction-order change, oracle-diff gated. Guarded * by QWEN35_SSM_2PASS (default OFF); the baseline kernel stays the bit-exact * serve path. Grid/block identical to fc_ssm_recurrence_slots. */ __global__ void fc_ssm_recurrence_slots_2pass(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, float scale) { extern __shared__ float sdata[]; /* [3*S]: pred | oq | qk partials */ const int S = state_size; const int j = blockIdx.x; /* row j of the transposed state */ const int h = blockIdx.y; /* V-head */ const int u = blockIdx.z; /* slot (user) */ const int i = threadIdx.x; /* state coordinate within the row */ if (i >= S) return; float *row = s + (size_t)u * n_v_heads * S * S + ((size_t)h * S + j) * S; const float *kh = k + (size_t)u * n_v_heads * S + (size_t)h * S; const float *qh = q + (size_t)u * n_v_heads * S + (size_t)h * S; const float decay = expf(g[(size_t)u * n_v_heads + h]); const float beta_h = beta[(size_t)u * n_v_heads + h]; /* Decay in place (per-element, same as baseline -> bit-identical row_dec). */ const float ri = row[i] * decay; const float ki = kh[i]; const float qi = qh[i]; /* ONE fused reduction: pred=sum ri*k, oq=sum ri*q, qk=sum k*q. */ float *p_pred = sdata; float *p_oq = sdata + S; float *p_qk = sdata + 2 * S; p_pred[i] = ri * ki; p_oq[i] = ri * qi; p_qk[i] = ki * qi; __syncthreads(); for (int stride = S / 2; stride > 0; stride >>= 1) { if (i < stride) { p_pred[i] += p_pred[i + stride]; p_oq[i] += p_oq[i + stride]; p_qk[i] += p_qk[i + stride]; } __syncthreads(); } const float pred = p_pred[0]; const float oq = p_oq[0]; const float qk = p_qk[0]; const float delta = (v[(size_t)u * n_v_heads * S + (size_t)h * S + j] - pred) * beta_h; /* rank-1 write (state carries forward, byte-layout identical). */ row[i] = ri + ki * delta; /* T4 output: y[j] = scale*(oq + qk*delta) — no second reduction. */ if (i == 0) y[(size_t)u * n_v_heads * S + (size_t)h * S + j] = scale * (oq + qk * delta); } /* Decay gate over n positions: g[t][h] = ssm_a[h] * softplus(alpha[t][h] + * ssm_dt[h]). One thread per (t, h); n_heads gives the per-position head * count for the broadcast vectors. total = n * n_heads. */ __global__ void fc_ssm_decay_gate_nx(const float *alpha, const float *ssm_dt, const float *ssm_a, float *g, int n_heads, int total) { int idx = blockIdx.x * blockDim.x + threadIdx.x; if (idx >= total) return; int h = idx % n_heads; float v = alpha[idx] + ssm_dt[h]; float sp = (v > 20.0f) ? v : log1pf(expf(v)); g[idx] = ssm_a[h] * sp; } /* q/k L2-norm + GQA repeat over n positions (grid.y = position), with the v * passthrough folded in (each block copies its val_dim / n_k_heads slice of * v — a plain copy, no fp math, replacing the launcher's DtoD). */ __global__ void fc_ssm_qknorm_gqa_nx(const float *conv_out, float *qg, float *kg, float *vg, int n_k_heads, int kv_groups, int S, int key_dim, int val_dim, float eps) { int kh = blockIdx.x; int t = blockIdx.y; if (kh >= n_k_heads) return; int i = threadIdx.x; extern __shared__ float sh[]; int conv_dim = 2 * key_dim + val_dim; const float *q_src = conv_out + (size_t)t * conv_dim; const float *k_src = q_src + key_dim; const float *v_src = q_src + 2 * key_dim; float *qg_t = qg + (size_t)t * val_dim; float *kg_t = kg + (size_t)t * val_dim; float *vg_t = vg + (size_t)t * val_dim; float qv = q_src[(size_t)kh * S + i]; float kv = k_src[(size_t)kh * S + i]; sh[i] = qv * qv; sh[blockDim.x + i] = kv * kv; __syncthreads(); for (int stride = blockDim.x >> 1; stride > 0; stride >>= 1) { if (i < stride) { sh[i] += sh[i + stride]; sh[blockDim.x + i] += sh[blockDim.x + i + stride]; } __syncthreads(); } float inv_q = 1.0f / fmaxf(sqrtf(sh[0]), eps); float inv_k = 1.0f / fmaxf(sqrtf(sh[blockDim.x]), eps); float qn = qv * inv_q; float kn = kv * inv_k; for (int r = 0; r < kv_groups; r++) { int vh = r * n_k_heads + kh; qg_t[(size_t)vh * S + i] = qn; kg_t[(size_t)vh * S + i] = kn; } /* v passthrough: this block copies slice kh of val_dim. */ int slice = val_dim / n_k_heads; for (int e = i; e < slice; e += blockDim.x) vg_t[(size_t)kh * slice + e] = v_src[(size_t)kh * slice + e]; } } /* extern "C" */