Skip to main content

GDN_KERNELS

Constant GDN_KERNELS 

Source
pub const GDN_KERNELS: &str = r#"
// Split one f32 activation into two int8 words: a ~= s*(qh + ql/127). Lives here because it is the first
// device code emitted, and both the GDN gate norm and the projection producers need it.
__device__ __forceinline__ void lgx_q8hl(float t, signed char* __restrict__ qh,
                                         signed char* __restrict__ ql, long i){
  int h = __float2int_rn(t);
  h = h > 127 ? 127 : (h < -127 ? -127 : h);
  int l = __float2int_rn((t - (float)h) * 127.0f);
  qh[i] = (signed char)h;
  ql[i] = (signed char)l;
}

__device__ __forceinline__ float gdn_warp_sum(float v){
  for(int off=16; off>0; off>>=1) v += __shfl_down_sync(0xffffffffu, v, off);
  return v;
}
// Three block-wide reductions sharing one traversal of the memory row.
__device__ __forceinline__ void gdn_block_sum3(float a, float b, float c, float* oa, float* ob, float* oc){
  __shared__ float sa[32], sb[32], sc[32];
  int lane = threadIdx.x & 31, warp = threadIdx.x >> 5, nw = (blockDim.x + 31) >> 5;
  a = gdn_warp_sum(a); b = gdn_warp_sum(b); c = gdn_warp_sum(c);
  if(lane == 0){ sa[warp] = a; sb[warp] = b; sc[warp] = c; }
  __syncthreads();
  if(threadIdx.x == 0){
    float ta = 0.f, tb = 0.f, tc = 0.f;
    for(int i = 0; i < nw; ++i){ ta += sa[i]; tb += sb[i]; tc += sc[i]; }
    *oa = ta; *ob = tb; *oc = tc;
  }
  __syncthreads();
}

// grid = (HV, D), block = D threads. Thread b owns column b of row a of head h's memory.
// ONE WARP PER STATE ROW. The state is D x D per value head; row `a` needs S[a]·k, S[a]·q and k·q, which is
// three reductions over D. Giving the row to a warp makes all three warp shuffles over registers the lanes
// already hold — no __syncthreads, and the row is re-used from those same registers for the write-back, so
// the D x D memory is read once and written once. The block-per-row form this replaces launched D blocks per
// head (295k blocks per token across 48 layers) and paid three block-wide barriers to reduce one float each.
// `sreg` is sized for D <= 256; the loops are lane-strided so any D, including D < 32, is handled.
__global__ void gdn_step(const float* __restrict__ q, const float* __restrict__ k, const float* __restrict__ v,
                         const float* __restrict__ bet, const float* __restrict__ alp,
                         float* __restrict__ S, float* __restrict__ o, int D, int NK){
  const unsigned m = 0xffffffffu;
  int h = blockIdx.x, lane = threadIdx.x & 31, warp = threadIdx.x >> 5;
  int a = blockIdx.y * (blockDim.x >> 5) + warp;
  if(a >= D) return;
  // The key head this value head reads. llama.cpp broadcasts k/q from NK heads to NV with
  // `ggml_repeat_4d`, and ggml_repeat TILES (index modulo) rather than repeat-interleaving — so the
  // pairing is h % NK ([0,1,..,15,0,1,..,15,..]), NOT h / (NV/NK) ([0,0,0,1,1,1,..]).
  int g = h % NK;
  const float* kk = k + (long)g * D;
  const float* qq = q + (long)g * D;
  long base = (long)h * D * D + (long)a * D;

  float sreg[8];
  float u_ = 0.f, w_ = 0.f, kq_ = 0.f;
  int n = 0;
  for(int b = lane; b < D; b += 32){
    float sv = S[base + b], kb = kk[b], qb = qq[b];
    sreg[n++] = sv; u_ += sv * kb; w_ += sv * qb; kq_ += kb * qb;
  }
  for(int sh = 16; sh > 0; sh >>= 1){
    u_ += __shfl_down_sync(m, u_, sh); w_ += __shfl_down_sync(m, w_, sh); kq_ += __shfl_down_sync(m, kq_, sh);
  }
  u_ = __shfl_sync(m, u_, 0); w_ = __shfl_sync(m, w_, 0); kq_ = __shfl_sync(m, kq_, 0);

  float alpha = alp[h], beta = bet[h];
  float coef  = beta * (v[(long)h * D + a] - alpha * u_);
  // The 1/sqrt(D) output scale llama.cpp applies (`attn_col * scale`). It does NOT cancel under the
  // following RMSNorm: the recurrence output is tiny (RMS ~1e-2), so `eps` is comparable to the scaled
  // variance and `s*x/sqrt(s^2*m + eps) != x/sqrt(m + eps)`. Dropping it scaled the whole layer by ~1.25x.
  if(lane == 0) o[(long)h * D + a] = (alpha * w_ + coef * kq_) * rsqrtf((float)D);
  n = 0;
  for(int b = lane; b < D; b += 32) S[base + b] = alpha * sreg[n++] + coef * kk[b];
}

// grid = (HV), block = D. RMSNorm the head and scale by the shared weight, THEN gate by silu(z) — the order
// llama.cpp's `build_norm_gated` uses for this architecture. (HF's Qwen3-Next gates first; Qwen3.5 does not,
// and the checkpoint follows Qwen3.5.)
__global__ void gdn_gate_norm(const float* __restrict__ z, const float* __restrict__ gn,
                              const float* __restrict__ o, float* __restrict__ out, int D, float eps){
  int h = blockIdx.x, a = threadIdx.x;
  long i = (long)h * D + a;
  float zz = z[i];
  float val = o[i];
  __shared__ float ss[32];
  int lane = threadIdx.x & 31, warp = threadIdx.x >> 5, nw = (blockDim.x + 31) >> 5;
  float p = gdn_warp_sum(val * val);
  if(lane == 0) ss[warp] = p;
  __syncthreads();
  __shared__ float den;
  if(threadIdx.x == 0){ float t = 0.f; for(int j = 0; j < nw; ++j) t += ss[j]; den = sqrtf(t / (float)D + eps); }
  __syncthreads();
  out[i] = (gn[a] * val / den) * (zz / (1.0f + __expf(-zz)));
}

// The same epilogue with the ACTIVATION QUANTIZATION fused in. The o-projection that follows consumes int8,
// and `blockDim == D` with `D % 32 == 0` means one warp owns exactly one 32-element quantization group — so
// the group amax is a warp shuffle over registers the block already holds, not a second pass over global
// memory. Saves a launch and a full round-trip of the value vector per GDN layer.
__global__ void gdn_gate_norm_q(const float* __restrict__ z, const float* __restrict__ gn,
                                const float* __restrict__ o, float* __restrict__ out,
                                signed char* __restrict__ xq, float* __restrict__ xs, int D, float eps){
  int h = blockIdx.x, a = threadIdx.x;
  long i = (long)h * D + a;
  float zz = z[i];
  float val = o[i];
  __shared__ float ss[32];
  int lane = threadIdx.x & 31, warp = threadIdx.x >> 5, nw = (blockDim.x + 31) >> 5;
  float p = gdn_warp_sum(val * val);
  if(lane == 0) ss[warp] = p;
  __syncthreads();
  __shared__ float den;
  if(threadIdx.x == 0){ float t = 0.f; for(int j = 0; j < nw; ++j) t += ss[j]; den = sqrtf(t / (float)D + eps); }
  __syncthreads();
  float y = (gn[a] * val / den) * (zz / (1.0f + __expf(-zz)));
  out[i] = y;
  float am = fabsf(y);
  for(int sh = 16; sh > 0; sh >>= 1) am = fmaxf(am, __shfl_down_sync(0xffffffffu, am, sh));
  am = __shfl_sync(0xffffffffu, am, 0);
  float sc = am > 0.f ? am / 127.0f : 1.0f;
  if(lane == 0) xs[i >> 5] = sc;
  xq[i] = (signed char)__float2int_rn(y / sc);
}

// The gate norm with the two-word activation split, for the lane that keeps the activation nearly exact.
__global__ void gdn_gate_norm_hl(const float* __restrict__ z, const float* __restrict__ gn,
                                const float* __restrict__ o, float* __restrict__ out,
                                signed char* __restrict__ xq, signed char* __restrict__ xl,
                                 float* __restrict__ xs, int D, float eps){
  int h = blockIdx.x, a = threadIdx.x;
  long i = (long)h * D + a;
  float zz = z[i];
  float val = o[i];
  __shared__ float ss[32];
  int lane = threadIdx.x & 31, warp = threadIdx.x >> 5, nw = (blockDim.x + 31) >> 5;
  float p = gdn_warp_sum(val * val);
  if(lane == 0) ss[warp] = p;
  __syncthreads();
  __shared__ float den;
  if(threadIdx.x == 0){ float t = 0.f; for(int j = 0; j < nw; ++j) t += ss[j]; den = sqrtf(t / (float)D + eps); }
  __syncthreads();
  float y = (gn[a] * val / den) * (zz / (1.0f + __expf(-zz)));
  out[i] = y;
  float am = fabsf(y);
  for(int sh = 16; sh > 0; sh >>= 1) am = fmaxf(am, __shfl_down_sync(0xffffffffu, am, sh));
  am = __shfl_sync(0xffffffffu, am, 0);
  float sc = am > 0.f ? am / 127.0f : 1.0f;
  if(lane == 0) xs[i >> 5] = sc;
  lgx_q8hl(y / sc, xq, xl, i);
}
"#;
Expand description

One block per (value head, row) computes that row’s u, w and the shared k·q in a single pass over the memory, then rewrites the row from the register it already holds. o lands in scratch; the gated RMSNorm needs the whole head, so it is a second (tiny) kernel.