Skip to main content

QWEN38_KERNELS

Constant QWEN38_KERNELS 

Source
pub const QWEN38_KERNELS: &str = r#"
// Depthwise causal conv over C channels with a `T`-deep ring, silu FUSED into the epilogue. The ring is the
// layer's carried state: shift up one, append this token, reduce. Depthwise means no cross-channel mixing, so
// one thread owns a channel end to end and the whole thing is a single pass.
__global__ void gk_conv_silu(const float* __restrict__ x, const float* __restrict__ wc,
                             float* __restrict__ cs, float* __restrict__ y, int C, int T){
  int c = blockIdx.x * blockDim.x + threadIdx.x; if(c >= C) return;
  for(int t = 0; t < T - 1; ++t) cs[(long)t * C + c] = cs[(long)(t + 1) * C + c];
  cs[(long)(T - 1) * C + c] = x[c];
  float acc = 0.f;
  for(int t = 0; t < T; ++t) acc += cs[(long)t * C + c] * wc[(long)t * C + c];
  y[c] = acc / (1.0f + __expf(-acc));            // silu, fused
}

// The conv with BOTH epilogues fused: silu, then — for the leading `nl2` head-blocks, which hold q and k —
// the per-head L2 normalize GDN applies before the delta update. Launched one block per head-sized chunk
// (`blockDim == D`, `grid == C / D`), so the normalize is a reduction over registers the block already owns.
// The v region (blocks >= nl2) takes the silu and stops. One launch where there were two, and q/k never
// leave the SM between the two steps.
__global__ void gk_conv_silu_l2(const float* __restrict__ x, const float* __restrict__ wc,
                                float* __restrict__ cs, float* __restrict__ y, int C, int T,
                                int nl2, float eps){
  int c = blockIdx.x * blockDim.x + threadIdx.x; if(c >= C) return;
  for(int t = 0; t < T - 1; ++t) cs[(long)t * C + c] = cs[(long)(t + 1) * C + c];
  cs[(long)(T - 1) * C + c] = x[c];
  float acc = 0.f;
  for(int t = 0; t < T; ++t) acc += cs[(long)t * C + c] * wc[(long)t * C + c];
  float v = acc / (1.0f + __expf(-acc));          // silu, fused
  if(blockIdx.x >= nl2){ y[c] = v; return; }      // the v region needs no normalize
  __shared__ float ss[32];
  int lane = threadIdx.x & 31, warp = threadIdx.x >> 5, nw = (blockDim.x + 31) >> 5;
  float p = gdn_warp_sum(v * v);
  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 + eps); }
  __syncthreads();
  y[c] = v / den;
}

// Per-head L2 normalize in place — GDN normalizes q and k before the delta update. grid = heads, block = D.
__global__ void gk_l2norm_heads(float* __restrict__ x, int D, float eps){
  int h = blockIdx.x, i = threadIdx.x; long o = (long)h * D + i;
  float v = x[o];
  __shared__ float ss[32];
  int lane = threadIdx.x & 31, warp = threadIdx.x >> 5, nw = (blockDim.x + 31) >> 5;
  float p = gdn_warp_sum(v * v);
  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 + eps); }
  __syncthreads();
  x[o] = v / den;
}

// Both GDN gates in one launch: write strength and forget rate, one value per value head.
//   beta  = sigmoid(br)
//   decay = (1 + e^(ar + dt))^rate
// The paper writes decay as exp(−exp(A_log)·softplus(ar+dt)), but the GGUF conversion ALREADY stores
// −exp(A_log) in `ssm_a` (verified: all 48 values negative, |v| < 0.08), so `rate` is used directly — 
// exponentiating it again turns a ^(−0.04) long memory into a ^(−0.96) one that forgets almost every token.
// softplus is ln(1+e^u), so the logarithm cancels against the outer exponential: one fewer transcendental.
__global__ void gk_gdn_gates(const float* __restrict__ br, const float* __restrict__ ar,
                             const float* __restrict__ dt, const float* __restrict__ alog,
                             float* __restrict__ beta, float* __restrict__ decay, int NV){
  int i = blockIdx.x * blockDim.x + threadIdx.x; if(i >= NV) return;
  beta[i] = 1.0f / (1.0f + __expf(-br[i]));
  decay[i] = powf(1.0f + __expf(ar[i] + dt[i]), alog[i]);
}

// Per-head RMSNorm + PARTIAL RoPE. Qwen3.8 rotates `rot` of each head's `hd` dims (64 of 256) and passes the
// rest through, and the inverse frequencies are computed over `rot`, not `hd` — using `hd` there would detune
// every rotation while still producing plausible-looking output.
__global__ void gk_qk_norm_rope_part(float* __restrict__ base, const int* __restrict__ d_pos, int stride,
                                     const float* __restrict__ nw, int n_heads, int hd, int rot,
                                     float rbase, float eps){
  int head = blockIdx.x; if(head >= n_heads) return;
  float* h = base + (long)(*d_pos) * stride + (long)head * hd;
  __shared__ float sh[128];
  float s = 0.f; for(int i = threadIdx.x; i < hd; i += blockDim.x) s += h[i] * h[i];
  sh[threadIdx.x] = s; __syncthreads();
  for(int k = blockDim.x / 2; k > 0; k >>= 1){ if(threadIdx.x < k) sh[threadIdx.x] += sh[threadIdx.x + k]; __syncthreads(); }
  float sc = rsqrtf(sh[0] / hd + eps); __syncthreads();
  for(int i = threadIdx.x; i < hd; i += blockDim.x) h[i] = h[i] * sc * nw[i];
  __syncthreads();
  int half = rot / 2, pos = *d_pos;
  for(int i = threadIdx.x; i < half; i += blockDim.x){
    float f = powf(rbase, -2.0f * (float)i / (float)rot);       // over the ROTARY width
    float a = pos * f, c = cosf(a), si = sinf(a);
    float x0 = h[i], x1 = h[i + half];
    h[i] = x0 * c - x1 * si; h[i + half] = x0 * si + x1 * c;
  }
  // dims [rot, hd) are carried through untouched — that is what "partial" means.
}

// Gated attention's output gate: ctx *= sigmoid(g). The checkpoint packs g into the Q projection, so this is
// the only extra work a gated attention layer does over an ordinary one.
__global__ void gk_gate_mul(float* __restrict__ ctx, const float* __restrict__ g, int n){
  int i = blockIdx.x * blockDim.x + threadIdx.x;
  if(i < n) ctx[i] *= 1.0f / (1.0f + __expf(-g[i]));
}
"#;
Expand description

The rest of the Qwen3.8-specific device code: the pieces of a GDN layer that surround the recurrence, plus the two things gated attention needs that a plain transformer does not (a partial rotary and an output gate). Each is deliberately fused past what a per-operator graph could do — the conv carries silu, the gate kernel computes both gates from one launch, and the norm+rope stay one pass over the head.