328 lines
12 KiB
Plaintext
328 lines
12 KiB
Plaintext
/* SPDX-License-Identifier: AGPL-3.0-or-later */
|
|
/* Copyright (C) 2026 David (Corey) Herrera */
|
|
/*
|
|
* moe_routing.cu — Qwen3.5-MoE routing kernels for Foundry
|
|
*
|
|
* Target: sm_37 (K80, Kepler)
|
|
* Compile: nvcc -ptx -arch=sm_37 moe_routing.cu -o moe_routing.ptx
|
|
*
|
|
* The MoE FFN splits into three GPU-side pieces. The heavy per-expert SwiGLU
|
|
* matmuls are NOT here — they are the existing quantized matvec kernels
|
|
* (fc_matvec_q3k + fc_silu + elementwise mul), launched per selected expert by
|
|
* the host. This file owns the cheap glue that surrounds them:
|
|
*
|
|
* 1. fc_moe_router_topk — softmax(router_logits) over n_expert, select
|
|
* top-k by probability, norm_topk_prob renormalize,
|
|
* apply expert_weights_scale. Emits top_idx[k] and
|
|
* top_wt[k] for the host to drive expert dispatch.
|
|
* 2. fc_moe_weighted_add — accumulate one expert's output into moe_out with
|
|
* its routed weight (the scatter/sum step).
|
|
* 3. fc_moe_shared_gate — sigmoid scalar gate applied to the shared-expert
|
|
* output (always-on branch).
|
|
*
|
|
* Mirrors, value-for-value, the CPU reference:
|
|
* Pantheon/src/models/qwen35moe_moe.c
|
|
* qwen35moe_router_softmax (lines 47-61)
|
|
* qwen35moe_select_topk (lines 65-89)
|
|
* norm_topk_prob + scale (lines 131-136)
|
|
* weighted expert sum (lines 138-147)
|
|
* shared-expert sigmoid gate (lines 179-185)
|
|
*
|
|
* Config: n_expert 256, n_expert_used 8, n_embd 2048, expert_ff 512.
|
|
*
|
|
* Head-broadcast trap (qwen35moe.md watch-list): MoE routing does NO GQA head
|
|
* broadcasting, so the interleaved-vs-blocked tiling bug does not apply in this
|
|
* file. Noted here so the integration step doesn't have to re-derive it.
|
|
*
|
|
* PTX symbols emitted (for embedding into ../cuda/fc_cuda_ptx.h):
|
|
* FC_PTX_MOE_ROUTING -> kernels
|
|
* fc_moe_router_topk, fc_moe_weighted_add, fc_moe_shared_gate
|
|
*
|
|
* Kepler notes: pure fp32; single-block reductions via shared memory (no
|
|
* __shfl). Top-k selection is a serial selection sort on thread 0 — k=8 and
|
|
* n_expert=256, so it is cheap and exactly reproduces the CPU tie-break order.
|
|
*/
|
|
|
|
extern "C" {
|
|
|
|
/* Hard cap on top-k — sized for the on-block selection scratch. Matches the
|
|
* CPU PN_QWEN35MOE_MAX_TOPK so behaviour is identical for any configured k. */
|
|
#define FC_MOE_MAX_TOPK 32
|
|
|
|
/*
|
|
* Router: softmax over n_expert logits, top-k select, norm_topk_prob, scale.
|
|
*
|
|
* logits: [n_expert] fp32 router logits (raw ffn_gate_inp . x). Overwritten
|
|
* in place with the softmax probabilities.
|
|
* top_idx: [k] int32 — selected expert indices (output)
|
|
* top_wt: [k] fp32 — final routed weights (output), already renormalized and
|
|
* multiplied by expert_weights_scale.
|
|
* n_expert: number of experts (256)
|
|
* k: n_expert_used (8); must be <= FC_MOE_MAX_TOPK
|
|
* expert_weights_scale: cfg->expert_weights_scale (1.0 unless GGUF overrides)
|
|
*
|
|
* Launched with 1 block, blockDim.x threads (use 256).
|
|
* Shared memory: blockDim.x * sizeof(float) * 2 (max reduction | sum reduction).
|
|
*
|
|
* Softmax + selection exactly follow the CPU path: selection sort picks the
|
|
* highest probability each round, skipping already-chosen experts, so ties
|
|
* resolve to the lowest index — identical to qwen35moe_select_topk.
|
|
*/
|
|
__global__ void fc_moe_router_topk(float *logits, int *top_idx, float *top_wt,
|
|
int n_expert, int k,
|
|
float expert_weights_scale) {
|
|
extern __shared__ float sdata[];
|
|
|
|
int tid = threadIdx.x;
|
|
int threads = blockDim.x;
|
|
float *smax = sdata;
|
|
float *ssum = sdata + threads;
|
|
|
|
/* Softmax pass 1: max */
|
|
float local_max = -1e30f;
|
|
for (int e = tid; e < n_expert; e += threads) {
|
|
float v = logits[e];
|
|
if (v > local_max) local_max = v;
|
|
}
|
|
smax[tid] = local_max;
|
|
__syncthreads();
|
|
for (int s = threads / 2; s > 0; s >>= 1) {
|
|
if (tid < s && smax[tid + s] > smax[tid])
|
|
smax[tid] = smax[tid + s];
|
|
__syncthreads();
|
|
}
|
|
float row_max = smax[0];
|
|
|
|
/* Softmax pass 2: exp(x - max) in place + sum */
|
|
float local_sum = 0.0f;
|
|
for (int e = tid; e < n_expert; e += threads) {
|
|
float v = expf(logits[e] - row_max);
|
|
logits[e] = v;
|
|
local_sum += v;
|
|
}
|
|
ssum[tid] = local_sum;
|
|
__syncthreads();
|
|
for (int s = threads / 2; s > 0; s >>= 1) {
|
|
if (tid < s)
|
|
ssum[tid] += ssum[tid + s];
|
|
__syncthreads();
|
|
}
|
|
float inv_sum = 1.0f / ssum[0];
|
|
|
|
/* Softmax pass 3: normalize to probabilities in place */
|
|
for (int e = tid; e < n_expert; e += threads)
|
|
logits[e] *= inv_sum;
|
|
__syncthreads();
|
|
|
|
/* Top-k via k rounds of PARALLEL argmax (tree reduction, ties -> lowest
|
|
* index, identical to the CPU first-seen-highest tie-break). The original
|
|
* serial selection sort on thread 0 measured 3.3 ms per launch (~134 ms
|
|
* per token across 40 layers) — one GPU thread scanning k*n_expert*k; the
|
|
* parallel version is microseconds. Selected entries are knocked out of
|
|
* logits[] in place with -1 (logits is per-token scratch; nothing reads
|
|
* the probabilities after this kernel). */
|
|
__shared__ float rv[256];
|
|
__shared__ int ri[256];
|
|
__shared__ float swt_sum;
|
|
if (tid == 0)
|
|
swt_sum = 0.0f;
|
|
|
|
for (int sel = 0; sel < k; sel++) {
|
|
/* Strided local argmax (ascending e + strict > => lowest index wins
|
|
* within a thread), then tree argmax preferring the lower index on
|
|
* equal values. */
|
|
float best_val = -1.0f;
|
|
int best = -1;
|
|
for (int e = tid; e < n_expert; e += threads) {
|
|
float v = logits[e];
|
|
if (v > best_val) {
|
|
best_val = v;
|
|
best = e;
|
|
}
|
|
}
|
|
rv[tid] = best_val;
|
|
ri[tid] = best;
|
|
__syncthreads();
|
|
for (int s = threads / 2; s > 0; s >>= 1) {
|
|
if (tid < s) {
|
|
if (rv[tid + s] > rv[tid] ||
|
|
(rv[tid + s] == rv[tid] && ri[tid + s] >= 0 &&
|
|
(ri[tid] < 0 || ri[tid + s] < ri[tid]))) {
|
|
rv[tid] = rv[tid + s];
|
|
ri[tid] = ri[tid + s];
|
|
}
|
|
}
|
|
__syncthreads();
|
|
}
|
|
if (tid == 0) {
|
|
top_idx[sel] = ri[0];
|
|
top_wt[sel] = rv[0];
|
|
swt_sum += rv[0];
|
|
if (ri[0] >= 0)
|
|
logits[ri[0]] = -1.0f; /* knock out for the next round */
|
|
}
|
|
__syncthreads();
|
|
}
|
|
|
|
/* norm_topk_prob: renormalize selected weights to sum to 1, then apply
|
|
* expert_weights_scale. */
|
|
if (tid == 0) {
|
|
float wt_sum = swt_sum;
|
|
float inv_wt = (wt_sum > 0.0f) ? (1.0f / wt_sum) : 0.0f;
|
|
float scale = expert_weights_scale * inv_wt;
|
|
for (int s = 0; s < k; s++)
|
|
top_wt[s] *= scale;
|
|
}
|
|
}
|
|
|
|
/*
|
|
* Batched (prefill) router: one token per BLOCK (blockIdx.x = t), each
|
|
* block running the exact fc_moe_router_topk body on its token's slice:
|
|
* logits [n_tok][n_expert] (softmax'd/knocked-out in place),
|
|
* top_idx [n_tok][k], top_wt [n_tok][k]. Same blockDim as the single-token
|
|
* launch -> same reduction trees -> bit-exact per token.
|
|
*/
|
|
__global__ void fc_moe_router_topk_nx(float *logits, int *top_idx,
|
|
float *top_wt, int n_expert, int k,
|
|
float expert_weights_scale) {
|
|
extern __shared__ float sdata[];
|
|
|
|
int tok = blockIdx.x;
|
|
float *lg = logits + (size_t)tok * n_expert;
|
|
int *ti = top_idx + (size_t)tok * k;
|
|
float *tw = top_wt + (size_t)tok * k;
|
|
|
|
int tid = threadIdx.x;
|
|
int threads = blockDim.x;
|
|
float *smax = sdata;
|
|
float *ssum = sdata + threads;
|
|
|
|
/* Softmax pass 1: max */
|
|
float local_max = -1e30f;
|
|
for (int e = tid; e < n_expert; e += threads) {
|
|
float v = lg[e];
|
|
if (v > local_max) local_max = v;
|
|
}
|
|
smax[tid] = local_max;
|
|
__syncthreads();
|
|
for (int s = threads / 2; s > 0; s >>= 1) {
|
|
if (tid < s && smax[tid + s] > smax[tid])
|
|
smax[tid] = smax[tid + s];
|
|
__syncthreads();
|
|
}
|
|
float row_max = smax[0];
|
|
|
|
/* Softmax pass 2: exp(x - max) in place + sum */
|
|
float local_sum = 0.0f;
|
|
for (int e = tid; e < n_expert; e += threads) {
|
|
float v = expf(lg[e] - row_max);
|
|
lg[e] = v;
|
|
local_sum += v;
|
|
}
|
|
ssum[tid] = local_sum;
|
|
__syncthreads();
|
|
for (int s = threads / 2; s > 0; s >>= 1) {
|
|
if (tid < s)
|
|
ssum[tid] += ssum[tid + s];
|
|
__syncthreads();
|
|
}
|
|
float inv_sum = 1.0f / ssum[0];
|
|
|
|
/* Softmax pass 3: normalize to probabilities in place */
|
|
for (int e = tid; e < n_expert; e += threads)
|
|
lg[e] *= inv_sum;
|
|
__syncthreads();
|
|
|
|
/* Top-k via k rounds of parallel argmax (ties -> lowest index). */
|
|
__shared__ float rv[256];
|
|
__shared__ int ri[256];
|
|
__shared__ float swt_sum;
|
|
if (tid == 0)
|
|
swt_sum = 0.0f;
|
|
|
|
for (int sel = 0; sel < k; sel++) {
|
|
float best_val = -1.0f;
|
|
int best = -1;
|
|
for (int e = tid; e < n_expert; e += threads) {
|
|
float v = lg[e];
|
|
if (v > best_val) {
|
|
best_val = v;
|
|
best = e;
|
|
}
|
|
}
|
|
rv[tid] = best_val;
|
|
ri[tid] = best;
|
|
__syncthreads();
|
|
for (int s = threads / 2; s > 0; s >>= 1) {
|
|
if (tid < s) {
|
|
if (rv[tid + s] > rv[tid] ||
|
|
(rv[tid + s] == rv[tid] && ri[tid + s] >= 0 &&
|
|
(ri[tid] < 0 || ri[tid + s] < ri[tid]))) {
|
|
rv[tid] = rv[tid + s];
|
|
ri[tid] = ri[tid + s];
|
|
}
|
|
}
|
|
__syncthreads();
|
|
}
|
|
if (tid == 0) {
|
|
ti[sel] = ri[0];
|
|
tw[sel] = rv[0];
|
|
swt_sum += rv[0];
|
|
if (ri[0] >= 0)
|
|
lg[ri[0]] = -1.0f;
|
|
}
|
|
__syncthreads();
|
|
}
|
|
|
|
/* norm_topk_prob + expert_weights_scale. */
|
|
if (tid == 0) {
|
|
float wt_sum = swt_sum;
|
|
float inv_wt = (wt_sum > 0.0f) ? (1.0f / wt_sum) : 0.0f;
|
|
float scale = expert_weights_scale * inv_wt;
|
|
for (int s = 0; s < k; s++)
|
|
tw[s] *= scale;
|
|
}
|
|
}
|
|
|
|
/*
|
|
* Weighted accumulate: moe_out[i] += weight * expert_out[i].
|
|
* The scatter/sum step of the routed-expert loop (qwen35moe_moe.c lines 144-146).
|
|
*
|
|
* expert_out: [n_embd] fp32 — one expert's SwiGLU output
|
|
* moe_out: [n_embd] fp32 — running accumulator (host zeroes it before s=0)
|
|
* weight: the routed weight top_wt[s] for this expert
|
|
*
|
|
* Launched with ceil(n_embd / blockDim.x) blocks, blockDim.x threads (use 256).
|
|
* No shared memory.
|
|
*/
|
|
__global__ void fc_moe_weighted_add(const float *expert_out, float *moe_out,
|
|
float weight, int n_embd) {
|
|
int i = blockIdx.x * blockDim.x + threadIdx.x;
|
|
if (i < n_embd)
|
|
moe_out[i] += weight * expert_out[i];
|
|
}
|
|
|
|
/*
|
|
* Shared-expert sigmoid gate: shexp_out[i] *= sigmoid(gate_logit).
|
|
* The scalar gate applied ONLY to the shared-expert branch
|
|
* (qwen35moe_moe.c lines 183-185). gate_logit = (ffn_gate_inp_shexp . x) is a
|
|
* single dot product computed by the host (the shared gate is a 1-D vector, so
|
|
* the CPU path computes it directly rather than via the matvec kernel).
|
|
*
|
|
* shexp_out: [n_embd] fp32 — shared-expert SwiGLU output, gated in place
|
|
* gate_logit: scalar pre-sigmoid logit
|
|
*
|
|
* Launched with ceil(n_embd / blockDim.x) blocks, blockDim.x threads (use 256).
|
|
* No shared memory.
|
|
*/
|
|
__global__ void fc_moe_shared_gate(float *shexp_out, float gate_logit,
|
|
int n_embd) {
|
|
int i = blockIdx.x * blockDim.x + threadIdx.x;
|
|
if (i < n_embd) {
|
|
float sgate = 1.0f / (1.0f + expf(-gate_logit));
|
|
shexp_out[i] *= sgate;
|
|
}
|
|
}
|
|
|
|
} /* extern "C" */
|