David
/
Foundry
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity
Foundry/cuda_kernels/moe_expert_group.cu

1623 lines
89 KiB
Plaintext
Raw Blame History

/* SPDX-License-Identifier: AGPL-3.0-or-later */
/* Copyright (C) 2026 David (Corey) Herrera */
/*
* moe_expert_group.cu — grouped top-k expert SwiGLU for Qwen3.5-MoE (Foundry)
*
* Target: sm_37 (K80, Kepler)
* Compile: nvcc -ptx -arch=sm_37 moe_expert_group.cu -o moe_expert_group.ptx
*
* Replaces the per-expert launch storm in the MoE FFN. The old path launched,
* per layer, 8 x (gate matvec + up matvec + silu + mul + down matvec +
* weighted_add) = 48 kernels plus a host sync + top-k read-back to drive the
* dispatch. These kernels do the whole routed-expert block in TWO launches,
* reading the router's top_idx/top_wt selection directly from device memory:
*
* 1. fc_moe_gateup_q3k — warp-per-row over (ff, k): computes
* hid[s*ff+row] = silu(gate[e_s].row . x) * (up[e_s].row . x)
* for e_s = top_idx[s]. Q3_K weights, both matrices contiguous 3D
* [n_expert, ff, cols] uploads (one fc_buffer per matrix per layer).
* 2. fc_moe_down_q3k / fc_moe_down_q4k — COOPERATIVE multi-warp-per-row-pair
* over n_embd: wpp warps split the k selected experts of one output row
* pair (warp wk takes slots wk, wk+wpp, ...), each accumulating weighted
* per-lane partials; one __shfl_down tree per warp plus a tiny shared-mem
* cross-warp sum then produce
* out[row] = sum_s top_wt[s] * (down[e_s].row . hid_s).
* Writes out directly (no zero + weighted_add passes). The 512-col down
* dot is only 2 super-blocks, so a lone warp per row-pair stalls on
* memory latency; wpp cooperating warps multiply the loads in flight.
* wpp is a launch arg — tuned default in fc_cuda.h (FC_MOE_DOWN_WPP).
*
* Memory paths (measured on a K80, nvprof): weight bytes use PLAIN
* global loads (L2, 32 B sectors) so the streamed-once weights stay off
* the read-only/tex cache. The down kernels read the tiny hid vector
* (16 KB, shared by every block) via __ldg from that now-private RO cache
* — this replaces the old per-slot shared staging, whose __syncthreads
* serialized the k-slot loop into latency-bound rounds (0.73 achieved
* occupancy -> 0.84 after). gateup keeps x staged in shared (each warp
* re-reads every x element 8x; LDS measured faster than __ldg there).
*
* Plus two small fusions for the always-on shared expert:
* 3. fc_moe_silu_mul — out = silu(a) * b in one pass (was 2 launches)
* 4. fc_moe_shexp_dot_gate — gate_logit = gate_w . x on-device, then
* shexp *= sigmoid(gate_logit)
*
* Weights are in the REPACKED device layouts produced by fc_repack_weight at
* upload (documented in fc_cuda.h and matvec_q3k.cu / matvec_kquants.cu):
* Q3_K and Q4_K 152 B aligned super-blocks laid out for the shift-free,
* cvt-free MAGIC-NUMBER dequant (sm_37 runs int-mul / int32->f32 cvt at
* 32/clk and shifts at 64/clk; the magic decode is AND + OR + one exact FMA
* per field — see matvec_q3k.cu for the full derivation). The lane dots
* below are duplicated from those files (PTX modules are standalone) and
* keep the identical per-lane fp op order. The gateup kernel stays
* bit-identical to the cvt-based kernels; the cooperative down kernels
* reassociate only the final slot-weighted reduction (per-warp slot subsets
* + cross-warp add) — every op is pinned with explicit __fmaf_rn/adds and
* reproduced exactly by the CPU replica in cuda/tests/bench_matvec.c, and
* the end-to-end argmax gate bounds the reassociation drift.
*
* QUANT COVERAGE. gate+up and down each come in a Q3_K and a Q4_K flavor, and
* the HOST dispatches both on the tensor's gguf type (fc_moe_gateup_ctx /
* fc_moe_down_accum_ctx, ../cuda/fc_cuda.c). Q3_K_M mixes them within one file
* (its merged ffn_down_exps is Q3_K on half the layers and Q4_K on the other
* half); a pure Q4_K_M model is Q4_K throughout, gate/up included, which is why
* fc_moe_gateup_q4k exists. DOWN additionally has a Q6_K flavor, because the
* Q4_K_M puts Q6_K on half the layers' merged ffn_down_exps. The DEDUP gate+up
* variants are Q3_K-only — see the dedup block comment below.
*
* PTX symbols emitted (for embedding into ../cuda/fc_cuda_ptx.h):
* FC_PTX_MOE_EXPERT_GROUP -> kernels fc_moe_gateup_q3k, fc_moe_gateup_q4k,
* fc_moe_down_q3k, fc_moe_down_q4k, fc_moe_down_q6k, fc_moe_silu_mul,
* fc_moe_shexp_dot_gate, plus the batched
* _nx variants and the cross-user dedup pair fc_moe_dedup_build /
* fc_moe_gateup_q3k_dedup (+ _shx, and + the compacted-grid pair
* fc_moe_dedup_build_wl / fc_moe_gateup_q3k_dedup_shxc in a
* -DFC_MOE_DEDUP_COMPACT=1 build only)
*
* The _shx / _shxc x-staging loop is selected by FC_MOE_SHX_STAGE
* (moe_dedup_shx_staging.cuh), default 0 = the shipped form; at 0 this module's
* PTX is byte-identical to the checked-in artifact.
*
* Kepler notes: pure fp32; __shfl_down is the non-_sync Kepler variant.
*/
extern "C" {
/* Repacked block sizes (bytes) — must match fc_repack_weight. */
#define FC_Q3K_RBLOCK_BYTES 152
#define FC_Q4K_RBLOCK_BYTES 152
#define FC_Q6K_RBLOCK_BYTES 216
/* Magic decodes (duplicated from matvec_q3k.cu / matvec_kquants.cu):
* exact float (v - bias) from the masked field at bit p. */
#define FC_MAGIC_NIB(w, p, bias) \
__fmaf_rn(__int_as_float(0x4B000000u | ((w) & (0xFu << (p)))), \
1.0f / (float)(1u << (p)), \
-(float)((1u << (23 - (p))) + (bias)))
#define FC_MAGIC_BYTE(w, p, bias) \
__fmaf_rn(__int_as_float(0x4B000000u | ((w) & (0xFFu << (p)))), \
1.0f / (float)(1u << (p)), \
-(float)((1u << (23 - (p))) + (bias)))
/* 2-bit field at bit p, scaled by 16 — Q6_K's high-bit pair. */
#define FC_MAGIC_BIT2X16(w, p) \
__fmaf_rn(__int_as_float(0x4B000000u | ((w) & (3u << (p)))), \
16.0f / (float)(1u << (p)), -(float)(1u << (27 - (p))))
/* Byte-aligned fields decode fastest through Kepler's 128/clk 8-bit cvt
* (single I2F.F32.U8 with byte select — cheaper than the magic FMA, which
* stays for the sub-byte fields). */
#define FC_CVT_BYTE(w, k) ((float)(((w) >> (8 * (k))) & 0xFFu))
__device__ __forceinline__ float fc_warp_sum_grp(float v) {
v += __shfl_down(v, 16);
v += __shfl_down(v, 8);
v += __shfl_down(v, 4);
v += __shfl_down(v, 2);
v += __shfl_down(v, 1);
return v;
}
/* x-load flavors for the lane dots below: the down kernels read the tiny
* hid vector straight from GLOBAL via __ldg (16 KB, read-only-cache
* resident), while the gateup kernel reads its block-staged x from SHARED
* (plain load — __ldg is global-only). */
#define FC_XLOAD_LDG(p, off) __ldg((p) + (off))
#define FC_XLOAD_PLAIN(p, off) (p)[off]
/*
* Interleaved ROW-PAIR lane dot: the two rows' 8-element FMA chains advance
* in lockstep per element (two independent chains the scheduler can
* dual-issue, unlike two back-to-back dependent 8-FMA chains), sharing one
* x load per element. Each row's per-lane chain keeps the exact i = 0..7
* order of fc_q3k_sb_lane_dot -> bit-exact per row.
* Q3_K repacked super-block layout (152 B): u32 qw[32] @0 (nibble i =
* qv | hbit<<2, value = nibble - 4), u8 sc[16] @128 (raw 6-bit, value =
* sc - 32, reordered per lane half), f32 d @144.
* Weights use PLAIN global loads (L2, 32 B sectors) — they stream through
* once, so keeping them off the read-only/tex path leaves that cache to
* the reused activation vector.
*/
#define FC_DEF_Q3K_GRP2(NAME, XLOAD) \
__device__ __forceinline__ void NAME(const unsigned char *blk0, \
const unsigned char *blk1, \
const float *xb, int lane, \
float *p0, float *p1) { \
unsigned int wv0 = ((const unsigned int *)blk0)[lane]; \
unsigned int wv1 = ((const unsigned int *)blk1)[lane]; \
float d0 = *(const float *)(blk0 + 144); \
float d1 = *(const float *)(blk1 + 144); \
uint2 sw0 = ((const uint2 *)(blk0 + 128))[lane >> 4]; \
uint2 sw1 = ((const uint2 *)(blk1 + 128))[lane >> 4]; \
unsigned int wh0 = wv0 >> 16, wh1 = wv1 >> 16; \
\
float a0 = 0.0f, a1 = 0.0f; \
_Pragma("unroll") \
for (int i = 0; i < 8; i++) { \
int qp = (i <= 4) ? 4 * i : 4 * i - 16; \
unsigned int q0 = (i <= 4) ? wv0 : wh0; \
unsigned int q1 = (i <= 4) ? wv1 : wh1; \
unsigned int s0 = (i < 4) ? sw0.x : sw0.y; \
unsigned int s1 = (i < 4) ? sw1.x : sw1.y; \
float xv = XLOAD(xb, 32 * i + lane); \
float fq0 = FC_MAGIC_NIB(q0, qp, 4); \
float fs0 = __fsub_rn(FC_CVT_BYTE(s0, i & 3), 32.0f); \
float fq1 = FC_MAGIC_NIB(q1, qp, 4); \
float fs1 = __fsub_rn(FC_CVT_BYTE(s1, i & 3), 32.0f); \
a0 = __fmaf_rn(__fmul_rn(fs0, fq0), xv, a0); \
a1 = __fmaf_rn(__fmul_rn(fs1, fq1), xv, a1); \
} \
*p0 += d0 * a0; \
*p1 += d1 * a1; \
}
FC_DEF_Q3K_GRP2(fc_q3k_sb_lane_grp2, FC_XLOAD_LDG) /* down: global hid */
FC_DEF_Q3K_GRP2(fc_q3k_sb_lane_grp2s, FC_XLOAD_PLAIN) /* gateup: shared x */
/* Interleaved ROW-PAIR Q4_K lane dot: both rows' chains advance per element
* (independent chains for dual-issue), sharing one x load. Each row keeps the
* exact per-element op order of fc_q4k_sb_lane_dot -> bit-exact. Q4_K repacked
* super-block layout (152 B): u32 qw[32] @0 (byte j = qs[32j+l]), u8 sc[8]
* @128, u8 mn[8] @136, f32 d @144, f32 dmin @148.
*
* Two XLOAD flavors, exactly like the Q3_K pair above: the down kernels read
* the tiny hid vector from GLOBAL via __ldg, the gate+up kernels read their
* block-staged x from SHARED. Made a macro (rather than a second copy) when
* the Q4_K gate+up kernels were added; the _grp2 instantiation expands to the
* same token stream the hand-written function had, so the down kernels' PTX is
* unchanged. */
#define FC_DEF_Q4K_GRP2(NAME, XLOAD) \
__device__ __forceinline__ void NAME(const unsigned char *blk0, \
const unsigned char *blk1, \
const float *xb, int lane, \
float *p0, float *p1) { \
/* weights: PLAIN global loads — see fc_q3k_sb_lane_grp2 */ \
unsigned int qw0 = ((const unsigned int *)blk0)[lane]; \
unsigned int qw1 = ((const unsigned int *)blk1)[lane]; \
uint2 sw0 = *(const uint2 *)(blk0 + 128); \
uint2 sw1 = *(const uint2 *)(blk1 + 128); \
uint2 mw0 = *(const uint2 *)(blk0 + 136); \
uint2 mw1 = *(const uint2 *)(blk1 + 136); \
float2 dd0 = *(const float2 *)(blk0 + 144); \
float2 dd1 = *(const float2 *)(blk1 + 144); \
unsigned int qh0 = qw0 >> 16, qh1 = qw1 >> 16; \
float a0 = 0.0f, a1 = 0.0f; \
_Pragma("unroll") \
for (int j = 0; j < 4; j++) { \
_Pragma("unroll") \
for (int h = 0; h < 2; h++) { \
int si = 2 * j + h; \
int nbit = 8 * j + 4 * h; \
int qp = (nbit <= 16) ? nbit : nbit - 16; \
unsigned int q0 = (nbit <= 16) ? qw0 : qh0; \
unsigned int q1 = (nbit <= 16) ? qw1 : qh1; \
unsigned int s0 = (si < 4) ? sw0.x : sw0.y; \
unsigned int s1 = (si < 4) ? sw1.x : sw1.y; \
unsigned int m0 = (si < 4) ? mw0.x : mw0.y; \
unsigned int m1 = (si < 4) ? mw1.x : mw1.y; \
float xv = XLOAD(xb, 64 * j + 32 * h + lane); \
float t0 = __fsub_rn( \
__fmul_rn(__fmul_rn(dd0.x, FC_CVT_BYTE(s0, si & 3)), \
FC_MAGIC_NIB(q0, qp, 0)), \
__fmul_rn(dd0.y, FC_CVT_BYTE(m0, si & 3))); \
float t1 = __fsub_rn( \
__fmul_rn(__fmul_rn(dd1.x, FC_CVT_BYTE(s1, si & 3)), \
FC_MAGIC_NIB(q1, qp, 0)), \
__fmul_rn(dd1.y, FC_CVT_BYTE(m1, si & 3))); \
a0 = __fmaf_rn(t0, xv, a0); \
a1 = __fmaf_rn(t1, xv, a1); \
} \
} \
*p0 += a0; \
*p1 += a1; \
}
FC_DEF_Q4K_GRP2(fc_q4k_sb_lane_grp2, FC_XLOAD_LDG) /* down: global hid */
FC_DEF_Q4K_GRP2(fc_q4k_sb_lane_grp2s, FC_XLOAD_PLAIN) /* gateup: shared x */
/* Interleaved ROW-PAIR Q6_K lane dot — duplicated from fused_layer.cu's
* fc_fu_q6k_acc2 / matvec_kquants.cu's fc_q6k_sb_lane_dot (PTX modules are
* standalone) with the identical per-element op order, so it is bit-exact with
* both. Q6_K repacked super-block layout (216 B): u32 ql[32] @0 (byte j =
* ql[32j+l]), u16 qh[32] @128, u8 sc[16] @192 (SIGNED, biased +128 in the
* repack), f32 d @208.
*
* Exists because the Q4_K_M quantization of a Qwen3.5-MoE puts Q6_K on half
* the layers' merged ffn_down_exps (the other half is Q4_K) — same shape of
* mixture as the Q3_K_M mini, one quant up. Without it those layers lose the
* grouped path entirely and fall back to 3*n_expert per-expert buffers a
* layer, which is what actually makes a Q4_K_M model unloadable on a K80 pair.
* Down-only: the gate+up matrices are never Q6_K in these files. */
#define FC_DEF_Q6K_GRP2(NAME, XLOAD) \
__device__ __forceinline__ void NAME(const unsigned char *blk0, \
const unsigned char *blk1, \
const float *xb, int lane, \
float *p0, float *p1) { \
unsigned int ql0 = ((const unsigned int *)blk0)[lane]; \
unsigned int ql1 = ((const unsigned int *)blk1)[lane]; \
unsigned int qh20 = ((const unsigned short *)(blk0 + 128))[lane]; \
unsigned int qh21 = ((const unsigned short *)(blk1 + 128))[lane]; \
float d0 = *(const float *)(blk0 + 208); \
float d1 = *(const float *)(blk1 + 208); \
uint2 sw0 = ((const uint2 *)(blk0 + 192))[lane >> 4]; \
uint2 sw1 = ((const uint2 *)(blk1 + 192))[lane >> 4]; \
unsigned int qlh0 = ql0 >> 16, qlh1 = ql1 >> 16; \
float a0 = 0.0f, a1 = 0.0f; \
_Pragma("unroll") \
for (int half = 0; half < 2; half++) { \
_Pragma("unroll") \
for (int quad = 0; quad < 4; quad++) { \
int i = 4 * half + quad; \
int nbit = 16 * half + 8 * (quad & 1) + 4 * ((quad >> 1) & 1); \
int hbit = 8 * half + 2 * quad; \
int qp = (nbit <= 16) ? nbit : nbit - 16; \
unsigned int q0 = (nbit <= 16) ? ql0 : qlh0; \
unsigned int q1 = (nbit <= 16) ? ql1 : qlh1; \
unsigned int s0 = (i < 4) ? sw0.x : sw0.y; \
unsigned int s1 = (i < 4) ? sw1.x : sw1.y; \
float xv = XLOAD(xb, 32 * i + lane); \
float fq0 = __fadd_rn(FC_MAGIC_NIB(q0, qp, 32), \
FC_MAGIC_BIT2X16(qh20, hbit)); \
float fq1 = __fadd_rn(FC_MAGIC_NIB(q1, qp, 32), \
FC_MAGIC_BIT2X16(qh21, hbit)); \
float fs0 = __fsub_rn(FC_CVT_BYTE(s0, i & 3), 128.0f); \
float fs1 = __fsub_rn(FC_CVT_BYTE(s1, i & 3), 128.0f); \
a0 = __fmaf_rn(__fmul_rn(fs0, fq0), xv, a0); \
a1 = __fmaf_rn(__fmul_rn(fs1, fq1), xv, a1); \
} \
} \
*p0 = __fmaf_rn(d0, a0, *p0); \
*p1 = __fmaf_rn(d1, a1, *p1); \
}
FC_DEF_Q6K_GRP2(fc_q6k_sb_lane_grp2, FC_XLOAD_LDG) /* down: global hid */
/*
* Interleaved ROW-PAIR F16 lane dot — gguf type 1, the TOP of this hardware's
* precision ladder (F32 storage does not fit: 35B x 4 B = 140 GB against 97.6
* GB of VRAM).
*
* sm_37 HAS NO FP16 ARITHMETIC — that starts at sm_53 — so this is not a copy
* of the fp32 path any more than it is of the quant ones: the weight is LOADED
* as 2 bytes and converted IN-REGISTER to fp32, and every multiply/add below is
* the same fp32 op the quant kernels use. cvt.f32.f16 is a STORAGE conversion
* and has existed since sm_20; it is what `__half2float` compiles to. Written
* as asm so this file stays header-free (no cuda_fp16.h), which is what keeps
* the checked-in PTX of every kernel above byte-identical when this block is
* added. Exact for subnormals and inf/nan — a widening cvt takes no .ftz.
*
* There is no block structure and no scale here, so an F16 "super-block" is
* just the same 256-element step the K-quant kernels already iterate in: 512 B
* of halves. FC_MOE_GATEUP_BODY / FC_MOE_DOWN_BODY are therefore reused
* verbatim with BLOCK_BYTES = 512 and row_stride = cols * 2. There is also no
* fc_repack_weight pass — the GGUF bytes ARE the device layout, which is what
* keeps a ~70 GB upload from needing a second host copy of every tensor on a
* 7.8 GB host.
*
* LOAD WIDTH. One 2-byte load per element at the same 32*i+lane lane->element
* map the K-quant dots use, so x staging, x indexing and the reduction tree are
* untouched. The warp's 32 lanes read 64 CONTIGUOUS bytes per i, so every load
* is fully coalesced and no DRAM byte is wasted: 16 32-B sectors per 512-B
* super-block, 100% efficiency.
*
* That is NOT the same as saying the load width does not matter, and an earlier
* version of this comment said it was — that a memory-bound kernel has idle
* issue slots either way. Coalescing settles DRAM bytes; it says nothing about
* how many LSU transactions those bytes cost. A warp-wide 16-bit load moves 64
* B in one transaction against a 128-B transaction budget, so the scalar form
* spends 8 half-empty transactions per 512-B super-block where a uint2 form
* spends 4 full ones. Measured: this kernel runs at 139 GB/s and the uint2 form
* fc_f16_sb_lane_grp2v below runs the same bytes at 181, so this map is leaving
* 1.3x on the table. See the FC_DEF_VEC16_GRP2 header for the numbers and for
* the lane-map change that is the price of them.
*/
#define FC_F16_RBLOCK_BYTES 512
__device__ __forceinline__ float fc_f16_to_f32(unsigned short h) {
float r;
asm("cvt.f32.f16 %0, %1;" : "=f"(r) : "h"(h));
return r;
}
#define FC_DEF_F16_GRP2(NAME, XLOAD) \
__device__ __forceinline__ void NAME(const unsigned char *blk0, \
const unsigned char *blk1, \
const float *xb, int lane, \
float *p0, float *p1) { \
/* weights: PLAIN global loads — see fc_q3k_sb_lane_grp2 */ \
const unsigned short *h0 = (const unsigned short *)blk0; \
const unsigned short *h1 = (const unsigned short *)blk1; \
float a0 = 0.0f, a1 = 0.0f; \
_Pragma("unroll") \
for (int i = 0; i < 8; i++) { \
int el = 32 * i + lane; \
float xv = XLOAD(xb, el); \
a0 = __fmaf_rn(fc_f16_to_f32(h0[el]), xv, a0); \
a1 = __fmaf_rn(fc_f16_to_f32(h1[el]), xv, a1); \
} \
*p0 += a0; \
*p1 += a1; \
}
FC_DEF_F16_GRP2(fc_f16_sb_lane_grp2, FC_XLOAD_LDG) /* down: global hid */
FC_DEF_F16_GRP2(fc_f16_sb_lane_grp2s, FC_XLOAD_PLAIN) /* gateup: shared x */
/*
* Interleaved ROW-PAIR BF16 lane dot — gguf type 30. Same 2-byte storage and
* same 512 B super-block as F16, and it exists for a reason F16 cannot serve:
* the model SHIPS as bf16 safetensors, so bf16 is the only format that carries
* the published weight VALUES. Converting bf16 -> F16 re-rounds every weight
* (bf16 keeps 7 stored mantissa bits, F16 keeps 10, and the two mantissa fields
* do not nest), so an F16 run measures a perturbed model. bf16 is therefore the
* control and F16 is the variant, not the other way round.
*
* IT IS NOT CHEAPER TO DECODE THAN F16, which is worth writing down because the
* arithmetic says it should be. A bf16 value IS the top 16 bits of the fp32 bit
* pattern, so widening is a shift and a free reinterpret, and the
* compute-capability 3.7 throughput table prices a 32-bit shift at 64 ops/clk/SM
* against 32 for "all other type conversions", which is where cvt.f32.f16 lands.
* Two arguments for bf16, and the SASS refutes both:
*
* fc_moe_gateup_f16 518 instructions, 16 F2F, 0 PRMT
* fc_moe_gateup_bf16 518 instructions, 0 F2F, 17 PRMT
*
* Same instruction count. ptxas implements the 16-bit shift as one PRMT (byte
* permute), not the two-op sequence the C suggests, so the widen is one
* instruction either way; and the throughput difference between the two classes
* never bills, because this kernel is memory-bound and its issue slots are not
* the constraint. Measured on die 7 at the per-layer shape: bf16 0.2400 ms vs
* f16 0.2418 ms on gate+up and 0.1300 vs 0.1304 on down — under 1%, inside the
* run-to-run spread, and it stays under 1% in the VEC form too even though the
* vec decodes differ in instruction count (see FC_DEF_VEC16_GRP2). So bf16 is
* here for FIDELITY — it is the only format holding the shipped values — and
* not for speed. Do not repeat the shift-beats-cvt argument; it has been
* measured, twice, and it does not reach the clock.
*
* Exactness: every bf16 is representable in fp32 with the mantissa zero-padded,
* including subnormals (a bf16 subnormal shifts to an fp32 subnormal with the
* same value — no renormalisation, which is the step F16 widening must do) and
* inf/nan (the payload rides along in the shifted mantissa). So the host replica
* is the identical shift and the compare is bit-exact for the dot.
*/
#define FC_BF16_RBLOCK_BYTES 512
__device__ __forceinline__ float fc_bf16_to_f32(unsigned short b) {
return __int_as_float((unsigned int)b << 16);
}
#define FC_DEF_BF16_GRP2(NAME, XLOAD) \
__device__ __forceinline__ void NAME(const unsigned char *blk0, \
const unsigned char *blk1, \
const float *xb, int lane, \
float *p0, float *p1) { \
/* weights: PLAIN global loads — see fc_q3k_sb_lane_grp2 */ \
const unsigned short *h0 = (const unsigned short *)blk0; \
const unsigned short *h1 = (const unsigned short *)blk1; \
float a0 = 0.0f, a1 = 0.0f; \
_Pragma("unroll") \
for (int i = 0; i < 8; i++) { \
int el = 32 * i + lane; \
float xv = XLOAD(xb, el); \
a0 = __fmaf_rn(fc_bf16_to_f32(h0[el]), xv, a0); \
a1 = __fmaf_rn(fc_bf16_to_f32(h1[el]), xv, a1); \
} \
*p0 += a0; \
*p1 += a1; \
}
FC_DEF_BF16_GRP2(fc_bf16_sb_lane_grp2, FC_XLOAD_LDG) /* down: global hid */
FC_DEF_BF16_GRP2(fc_bf16_sb_lane_grp2s, FC_XLOAD_PLAIN) /* gateup: shared x */
/*
* VECTORIZED (uint2) 16-bit lane dot — and it is worth 1.3x.
*
* WHY IT PAYS. The scalar dots above are perfectly coalesced but each warp-wide
* 16-bit load moves only 64 B, and an LSU transaction can carry 128 B. Eight
* such loads per 512-B super-block therefore cost 8 transactions at 50%
* occupancy. Loading a uint2 (four 16-bit weights, 8 B/lane) makes the warp read
* 256 CONTIGUOUS bytes per instruction = 2 full transactions, so a super-block
* costs 2 instructions / 4 transactions — the same DRAM bytes through half the
* LSU transactions. uint4 would issue 1 instruction for the same 4 transactions
* (no further transaction saving) while forcing an 8-float-per-lane stride on
* the staged x, which is an 8-way shared-memory bank conflict on Kepler's 4-byte
* banks; uint2 pairs with a float4 x load, the conflict-free pattern. So uint2
* is the widest useful width here, not a step on the way to uint4.
*
* MEASURED, die 7, the per-layer shape (k=8, ff=512, cols=2048), median of
* three runs that agreed to 0.3%, with the k experts forced DISTINCT so that
* nominal bytes and DRAM bytes are the same number:
*
* gateup f16 0.2418 ms 138.8 GB/s -> f16-vec 0.1849 ms 181.4 GB/s
* gateup bf16 0.2400 ms 139.8 GB/s -> bf16-vec 0.1844 ms 181.9 GB/s
* down f16 0.1304 ms 128.6 GB/s -> f16-vec 0.0986 ms 170.2 GB/s
* down bf16 0.1300 ms 129.1 GB/s -> bf16-vec 0.0980 ms 171.2 GB/s
*
* 1.30-1.33x on all four, and 182 GB/s is 76% of a K80 die's 240 GB/s — the
* band a well-tuned streaming kernel reaches on Kepler. The scalar form's 58%
* was the half-empty transactions and nothing else. The SASS confirms the
* mechanism rather than inferring it — LD.E.U16 x16 + LDS x8 becomes LD.E.64
* x12 + LDS.128 x6 for gateup, and LD.E.U16 x16 + LDG.E.T becomes LD.E.64 x12 +
* LDG.E.128 for down. It is paid for in registers, but not enough to matter:
* gateup 48 -> 52 and down 50 -> 56 (the _nx forms 53 -> 64), all under the 63
* that would cost a block of occupancy at 256 threads, and 0 bytes of stack
* frame either way — read the STACK FRAME line, not the spill counters, which
* hide local arrays.
*
* This is the first kernel in this tree where a load-side change converts to
* wall clock. The standing findings that the K80 is launch-bound and that work
* reduction does not convert were both learned from the K-quant kernels, which
* are instruction-bound; they do not transfer to a kernel whose bottleneck is
* the memory pipe, and this is the counter-example.
*
* The two formats stay within 1% of each other in the vec form as well, even
* though the vec decodes are no longer the same shape: F2F.F32.F16 carries a
* half-word selector so an F16 pair is 2 F2F, while a bf16 pair is a mask plus
* a permute — 48 F2F for f16-vec against 39 LOP32I + 25 PRMT for bf16-vec, at
* the same 595 instructions total. Even at 1.3x the kernel is still short of
* the memory pipe, so the decode difference does not reach the clock.
*
* LANE MAP CONSEQUENCE — this is not a drop-in. The scalar map gives lane l the
* strided set {32i + l}, which is exactly why it cannot be vectorized: those 8
* halves are not contiguous. The vec map gives lane l the contiguous runs
* {4l..4l+3} and {128 + 4l..128 + 4l+3}. Same 256 products per warp, different
* PARTITION of them across lanes, so the per-lane fmaf chains group differently
* and the shfl tree sums a different association. The result is therefore NOT
* bit-identical to the scalar kernel and must not be compared against it — the
* test carries a second replica built on the vec map and compares bit-exactly
* against THAT, which is what makes the comparison able to fail.
*
*/
#define FC_XLOAD4_LDG(p, i) __ldg((const float4 *)(p) + (i))
#define FC_XLOAD4_PLAIN(p, i) (((const float4 *)(p))[i])
/* Element pair packed little-endian in one 32-bit word: LO is the even element
* (low half), HI the odd one. The BF16 pair needs no conversion instruction —
* masking off the low half IS the fp32 value of the odd element, and shifting
* the low half up IS the even one — but that costs no fewer instructions than
* the F16 pair, because F2F.F32.F16 carries its own half-word selector and so
* reads either half of the word directly. Measured, both ways: see the
* fc_bf16_to_f32 header. */
#define FC_F16_LO(w) fc_f16_to_f32((unsigned short)((w) & 0xFFFFu))
#define FC_F16_HI(w) fc_f16_to_f32((unsigned short)((w) >> 16))
#define FC_BF16_LO(w) __int_as_float((w) << 16)
#define FC_BF16_HI(w) __int_as_float((w) & 0xFFFF0000u)
#define FC_DEF_VEC16_GRP2(NAME, CVT_LO, CVT_HI, XLOAD4) \
__device__ __forceinline__ void NAME(const unsigned char *blk0, \
const unsigned char *blk1, \
const float *xb, int lane, \
float *p0, float *p1) { \
float a0 = 0.0f, a1 = 0.0f; \
_Pragma("unroll") \
for (int h = 0; h < 2; h++) { \
/* 8 B/lane x 32 lanes = 256 contiguous bytes per warp instruction */ \
uint2 w0 = ((const uint2 *)(blk0 + 256 * h))[lane]; \
uint2 w1 = ((const uint2 *)(blk1 + 256 * h))[lane]; \
float4 xv = XLOAD4(xb + 128 * h, lane); \
/* the two rows' chains stay interleaved for dual-issue, exactly as \
* in the scalar dots; each row's own order is element-ascending */ \
a0 = __fmaf_rn(CVT_LO(w0.x), xv.x, a0); \
a1 = __fmaf_rn(CVT_LO(w1.x), xv.x, a1); \
a0 = __fmaf_rn(CVT_HI(w0.x), xv.y, a0); \
a1 = __fmaf_rn(CVT_HI(w1.x), xv.y, a1); \
a0 = __fmaf_rn(CVT_LO(w0.y), xv.z, a0); \
a1 = __fmaf_rn(CVT_LO(w1.y), xv.z, a1); \
a0 = __fmaf_rn(CVT_HI(w0.y), xv.w, a0); \
a1 = __fmaf_rn(CVT_HI(w1.y), xv.w, a1); \
} \
*p0 += a0; \
*p1 += a1; \
}
FC_DEF_VEC16_GRP2(fc_f16_sb_lane_grp2v, FC_F16_LO, FC_F16_HI,
FC_XLOAD4_LDG) /* down: global hid */
FC_DEF_VEC16_GRP2(fc_f16_sb_lane_grp2vs, FC_F16_LO, FC_F16_HI,
FC_XLOAD4_PLAIN) /* gateup: shared x */
FC_DEF_VEC16_GRP2(fc_bf16_sb_lane_grp2v, FC_BF16_LO, FC_BF16_HI,
FC_XLOAD4_LDG) /* down: global hid */
FC_DEF_VEC16_GRP2(fc_bf16_sb_lane_grp2vs, FC_BF16_LO, FC_BF16_HI,
FC_XLOAD4_PLAIN) /* gateup: shared x */
/*
* Grouped gate+up SwiGLU over the k selected Q3_K experts.
*
* gate_all/up_all: [n_expert * ff * row_stride] REPACKED Q3_K bytes — the
* merged 3D GGUF tensor, one contiguous upload per matrix.
* x: [cols] fp32 activation (the post_attention_norm'd hidden)
* top_idx: [k] int32 selected expert ids (device — no host read-back)
* hid: [k * ff] fp32 output, hid[s*ff + row]
* cols: inner dim (n_embd, 2048); row_stride = (cols/256)*152
*
* WARP-PER-ROW: blockDim = 256 (8 warps); warp w handles expert-FF row
* blockIdx.x * 8 + w of slot blockIdx.y, computing the gate AND up dots
* together against the shared-staged x (staged in 2048-element chunks;
* measured faster here than __ldg x — this kernel re-reads each x element
* 8x per warp and LDS keeps that off the tex path) — the two independent
* accumulator chains give this kernel its dual-issue ILP — then lane 0
* writes silu(g) * u. Inactive tail warps still hit the staging barriers.
* Launch: grid (ceil(ff/8), k), block 256,
* shared min(cols, 2048) * sizeof(float).
*/
/*
* The gate+up work itself, shared by the Q3_K and Q4_K entry points (and by
* their batched _nx forms, which differ only in XPTR/HIDOFF). A macro rather
* than an inline function for the same reason FC_MOE_DEDUP_SHX_BODY is one: it
* expands to the token stream the hand-written Q3_K kernel had, so that
* kernel's PTX is untouched by the Q4_K addition.
*
* Requires in scope: gate_all, up_all, hid, ff, cols, row_stride (parameters);
* sh_x, warp, lane, row, active, e (block-local prologue).
* XPTR — the token's x base pointer (x, or x + tok*cols for _nx)
* HIDOFF — the element index this warp writes in hid
* BLOCK_BYTES / SB_LANE_DOT2S — the quant's repacked stride and shared-x
* row-pair lane dot.
*/
#define FC_MOE_GATEUP_BODY(XPTR, HIDOFF, BLOCK_BYTES, SB_LANE_DOT2S) \
size_t expert_stride = (size_t)ff * (size_t)row_stride; \
/* strength-reduced running block pointers (one += per super-block) */ \
const unsigned char *gb = gate_all + (size_t)e * expert_stride + \
(size_t)row * (size_t)row_stride; \
const unsigned char *ub = up_all + (size_t)e * expert_stride + \
(size_t)row * (size_t)row_stride; \
\
float pg = 0.0f, pu = 0.0f; \
for (int cbase = 0; cbase < cols; cbase += 2048) { \
int cend = cols - cbase < 2048 ? cols - cbase : 2048; \
__syncthreads(); \
for (int i = threadIdx.x; i < cend; i += 256) \
sh_x[i] = __ldg((XPTR) + cbase + i); \
__syncthreads(); \
if (active) { \
int nsb = cend >> 8; \
const float *xb = sh_x; \
for (int sb = 0; sb < nsb; sb++) { \
/* gate/up element-interleaved: two independent chains */ \
SB_LANE_DOT2S(gb, ub, xb, lane, &pg, &pu); \
gb += BLOCK_BYTES; \
ub += BLOCK_BYTES; \
xb += 256; \
} \
} \
} \
\
pg = fc_warp_sum_grp(pg); \
pu = fc_warp_sum_grp(pu); \
if (active && lane == 0) \
hid[HIDOFF] = (pg / (1.0f + expf(-pg))) * pu;
/* Prologue shared by the two single-token gate+up entry points. */
#define FC_MOE_GATEUP_PROLOGUE \
int warp = threadIdx.x >> 5; \
int lane = threadIdx.x & 31; \
int row = blockIdx.x * 8 + warp; \
int slot = blockIdx.y; \
int active = row < ff; \
int e = __ldg(top_idx + slot);
__global__ void fc_moe_gateup_q3k(const unsigned char *gate_all,
const unsigned char *up_all,
const float *x, const int *top_idx,
float *hid, int ff, int cols,
int row_stride) {
extern __shared__ float sh_x[]; /* [min(cols, 2048)] staged x */
FC_MOE_GATEUP_PROLOGUE
FC_MOE_GATEUP_BODY(x, (size_t)slot * ff + row, FC_Q3K_RBLOCK_BYTES,
fc_q3k_sb_lane_grp2s)
}
/*
* Q4_K gate+up. Same kernel, same launch shape, same shared-x staging as
* fc_moe_gateup_q3k — only the super-block decode differs, and Q4_K's repacked
* super-block is the same 152 B, so cols/row_stride/grid arithmetic is
* unchanged. Exists so a pure Q4_K_M MoE (whose ffn_gate_exps/ffn_up_exps are
* Q4_K, not just the Q3_K_M mini's half-Q4_K ffn_down_exps) has a gate+up
* kernel at all; without it the host would feed Q4_K bytes to the Q3_K decode.
* Accumulates in fp32 like every other kernel here — sm_37 has no fp16 math.
*/
__global__ void fc_moe_gateup_q4k(const unsigned char *gate_all,
const unsigned char *up_all,
const float *x, const int *top_idx,
float *hid, int ff, int cols,
int row_stride) {
extern __shared__ float sh_x[]; /* [min(cols, 2048)] staged x */
FC_MOE_GATEUP_PROLOGUE
FC_MOE_GATEUP_BODY(x, (size_t)slot * ff + row, FC_Q4K_RBLOCK_BYTES,
fc_q4k_sb_lane_grp2s)
}
/*
* F16 gate+up. Same kernel, same launch shape, same shared-x staging as its
* Q3_K/Q4_K siblings — only the super-block stride (512 B, not 152) and the
* decode differ, and neither is in the grid arithmetic, which depends on
* ff/k/cols and not on row_stride. Accumulates in fp32 because sm_37 has no
* fp16 math (see fc_f16_to_f32).
*/
__global__ void fc_moe_gateup_f16(const unsigned char *gate_all,
const unsigned char *up_all,
const float *x, const int *top_idx,
float *hid, int ff, int cols,
int row_stride) {
extern __shared__ float sh_x[]; /* [min(cols, 2048)] staged x */
FC_MOE_GATEUP_PROLOGUE
FC_MOE_GATEUP_BODY(x, (size_t)slot * ff + row, FC_F16_RBLOCK_BYTES,
fc_f16_sb_lane_grp2s)
}
/*
* BF16 gate+up — fc_moe_gateup_f16 with the shift-widen decode. Identical
* launch shape and identical 512 B super-block; only the lane dot differs.
*/
__global__ void fc_moe_gateup_bf16(const unsigned char *gate_all,
const unsigned char *up_all,
const float *x, const int *top_idx,
float *hid, int ff, int cols,
int row_stride) {
extern __shared__ float sh_x[]; /* [min(cols, 2048)] staged x */
FC_MOE_GATEUP_PROLOGUE
FC_MOE_GATEUP_BODY(x, (size_t)slot * ff + row, FC_BF16_RBLOCK_BYTES,
fc_bf16_sb_lane_grp2s)
}
/*
* The uint2 gate+up pair. Same launch shape and the same staged x — the vec
* lane dot reads that x back as float4, so the staging array is declared
* float4 to pin the dynamic-shared base to a 16-B boundary (every offset
* inside it, 256 floats per super-block and 128 per half, is already a
* multiple of 16 B). Their hid is NOT bit-identical to the scalar kernels'
* (different lane->element partition — see FC_DEF_VEC16_GRP2).
*/
__global__ void fc_moe_gateup_f16_vec(const unsigned char *gate_all,
const unsigned char *up_all,
const float *x, const int *top_idx,
float *hid, int ff, int cols,
int row_stride) {
extern __shared__ float4 sh_x4[];
float *sh_x = (float *)sh_x4;
FC_MOE_GATEUP_PROLOGUE
FC_MOE_GATEUP_BODY(x, (size_t)slot * ff + row, FC_F16_RBLOCK_BYTES,
fc_f16_sb_lane_grp2vs)
}
__global__ void fc_moe_gateup_bf16_vec(const unsigned char *gate_all,
const unsigned char *up_all,
const float *x, const int *top_idx,
float *hid, int ff, int cols,
int row_stride) {
extern __shared__ float4 sh_x4[];
float *sh_x = (float *)sh_x4;
FC_MOE_GATEUP_PROLOGUE
FC_MOE_GATEUP_BODY(x, (size_t)slot * ff + row, FC_BF16_RBLOCK_BYTES,
fc_bf16_sb_lane_grp2vs)
}
/*
* Grouped down-projection + weighted accumulate over the k selected experts.
*
* down_all: [n_expert * rows * row_stride] REPACKED quant bytes (merged 3D)
* hid: [k * ff] fp32 (from fc_moe_gateup_q3k)
* top_idx: [k] int32, top_wt: [k] fp32 (device — no host read-back)
* out: [rows] fp32, WRITTEN (not accumulated into)
* ff: inner dim (expert_ff, 512); rows = n_embd
*
* COOPERATIVE MULTI-WARP-PER-ROW-PAIR: the block's warps are grouped wpp
* (warps-per-row-pair) at a time; group p handles output rows
* (blockIdx.x * n_warps/wpp + p) * 2 + {0,1} (row-pair: independent
* accumulators for dual-issue ILP, one hid load amortized across the pair).
* Within a group, warp wk takes slots wk, wk+wpp, ... of the k selected
* experts, chaining t = fma(top_wt[s], dot_s, t) per lane over ITS slots
* (hid read via __ldg — 16 KB total, read-only-cache resident, no staging
* barriers, so each warp's k/wpp short 2-super-block dots stream back to
* back with wpp x the loads in flight of the old one-warp-per-pair form).
* Each warp then tree-reduces its two lanesums (__shfl_down) and the wpp
* per-warp partials meet in a tiny static-shared cross-warp add, summed in
* ascending wk order by one thread per row (deterministic; the only fp
* reassociation vs. the old kernel — bounded by the argmax gate).
* wpp must divide the block's warp count; k need not divide wpp (idle-slot
* warps contribute 0).
* Launch: grid ceil(rows / (2 * n_warps/wpp)), block <= 512 (tuned: 128
* threads, wpp=2 — the kernel is instruction-issue bound at ~4.3 IPC / 78%
* issue-slot utilization, so beyond that extra warps buy nothing), shared 0.
*/
#define FC_MOE_DOWN_BODY(BLOCK_BYTES, SB_LANE_DOT2) \
__shared__ float sred[16][2]; /* per-warp row-pair partials, <=16 w */ \
int warp = threadIdx.x >> 5; \
int lane = threadIdx.x & 31; \
int pairs_per_block = (blockDim.x >> 5) / wpp; \
int pair = warp / wpp; \
int wk = warp - pair * wpp; \
int row0 = (blockIdx.x * pairs_per_block + pair) * 2; \
int active0 = row0 < rows; \
int active1 = row0 + 1 < rows; \
size_t expert_stride = (size_t)rows * (size_t)row_stride; \
int n_sb = ff >> 8; \
float t0 = 0.0f, t1 = 0.0f; \
if (active0) { \
for (int s = wk; s < k; s += wpp) { \
int e = __ldg(top_idx + s); \
const unsigned char *b0 = down_all + (size_t)e * expert_stride + \
(size_t)row0 * (size_t)row_stride; \
/* tail: clamp the pair onto row0, duplicate discarded */ \
const unsigned char *b1 = active1 ? b0 + (size_t)row_stride : b0; \
const float *hb = hid + (size_t)s * ff; \
float p0 = 0.0f, p1 = 0.0f; \
for (int sb = 0; sb < n_sb; sb++) { \
SB_LANE_DOT2(b0, b1, hb, lane, &p0, &p1); \
b0 += BLOCK_BYTES; \
b1 += BLOCK_BYTES; \
hb += 256; \
} \
float w = __ldg(top_wt + s); \
t0 = __fmaf_rn(w, p0, t0); \
t1 = __fmaf_rn(w, p1, t1); \
} \
} \
t0 = fc_warp_sum_grp(t0); \
t1 = fc_warp_sum_grp(t1); \
if (lane == 0) { \
sred[warp][0] = t0; \
sred[warp][1] = t1; \
} \
__syncthreads(); \
if ((int)threadIdx.x < 2 * pairs_per_block) { \
int p = threadIdx.x >> 1; \
int r = threadIdx.x & 1; \
int row = (blockIdx.x * pairs_per_block + p) * 2 + r; \
if (row < rows) { \
float acc = 0.0f; \
for (int w = 0; w < wpp; w++) \
acc += sred[p * wpp + w][r]; \
out[row] = acc; \
} \
}
__global__ void fc_moe_down_q3k(const unsigned char *down_all, const float *hid,
const int *top_idx, const float *top_wt,
float *out, int k, int rows, int ff,
int row_stride, int wpp) {
FC_MOE_DOWN_BODY(FC_Q3K_RBLOCK_BYTES, fc_q3k_sb_lane_grp2)
}
__global__ void fc_moe_down_q4k(const unsigned char *down_all, const float *hid,
const int *top_idx, const float *top_wt,
float *out, int k, int rows, int ff,
int row_stride, int wpp) {
FC_MOE_DOWN_BODY(FC_Q4K_RBLOCK_BYTES, fc_q4k_sb_lane_grp2)
}
__global__ void fc_moe_down_q6k(const unsigned char *down_all, const float *hid,
const int *top_idx, const float *top_wt,
float *out, int k, int rows, int ff,
int row_stride, int wpp) {
FC_MOE_DOWN_BODY(FC_Q6K_RBLOCK_BYTES, fc_q6k_sb_lane_grp2)
}
__global__ void fc_moe_down_f16(const unsigned char *down_all, const float *hid,
const int *top_idx, const float *top_wt,
float *out, int k, int rows, int ff,
int row_stride, int wpp) {
FC_MOE_DOWN_BODY(FC_F16_RBLOCK_BYTES, fc_f16_sb_lane_grp2)
}
__global__ void fc_moe_down_bf16(const unsigned char *down_all,
const float *hid, const int *top_idx,
const float *top_wt, float *out, int k,
int rows, int ff, int row_stride, int wpp) {
FC_MOE_DOWN_BODY(FC_BF16_RBLOCK_BYTES, fc_bf16_sb_lane_grp2)
}
/* uint2 down pair. hid comes from global here, so the float4 x load is an
* __ldg — hid is cudaMalloc'd (256-B aligned) and every offset into it
* (s * ff, 256 per super-block, 128 per half) is a multiple of 16 B. */
__global__ void fc_moe_down_f16_vec(const unsigned char *down_all,
const float *hid, const int *top_idx,
const float *top_wt, float *out, int k,
int rows, int ff, int row_stride,
int wpp) {
FC_MOE_DOWN_BODY(FC_F16_RBLOCK_BYTES, fc_f16_sb_lane_grp2v)
}
__global__ void fc_moe_down_bf16_vec(const unsigned char *down_all,
const float *hid, const int *top_idx,
const float *top_wt, float *out, int k,
int rows, int ff, int row_stride,
int wpp) {
FC_MOE_DOWN_BODY(FC_BF16_RBLOCK_BYTES, fc_bf16_sb_lane_grp2v)
}
/*
* Batched (prefill) grouped gate+up SwiGLU: grid.z = token t, everything
* else exactly fc_moe_gateup_q3k on token t's slices — x [n][cols],
* top_idx [n][k], hid [n][k*ff]. Per token bit-exact (same warp body).
* Weight reuse across tokens comes from L2 when tokens share experts.
*/
#define FC_MOE_GATEUP_NX_PROLOGUE \
int warp = threadIdx.x >> 5; \
int lane = threadIdx.x & 31; \
int row = blockIdx.x * 8 + warp; \
int slot = blockIdx.y; \
int tok = blockIdx.z; \
int active = row < ff; \
const float *xt = x + (size_t)tok * cols; \
int e = __ldg(top_idx + (size_t)tok * k + slot);
__global__ void fc_moe_gateup_q3k_nx(const unsigned char *gate_all,
const unsigned char *up_all,
const float *x, const int *top_idx,
float *hid, int k, int ff, int cols,
int row_stride) {
extern __shared__ float sh_x[];
FC_MOE_GATEUP_NX_PROLOGUE
FC_MOE_GATEUP_BODY(xt, ((size_t)tok * k + slot) * ff + row,
FC_Q3K_RBLOCK_BYTES, fc_q3k_sb_lane_grp2s)
}
/* Batched (prefill / multi-user decode) Q4_K gate+up — fc_moe_gateup_q4k with
* grid.z = token, exactly as _q3k_nx is to _q3k. Per token bit-exact with the
* single-token Q4_K kernel (same warp body). */
__global__ void fc_moe_gateup_q4k_nx(const unsigned char *gate_all,
const unsigned char *up_all,
const float *x, const int *top_idx,
float *hid, int k, int ff, int cols,
int row_stride) {
extern __shared__ float sh_x[];
FC_MOE_GATEUP_NX_PROLOGUE
FC_MOE_GATEUP_BODY(xt, ((size_t)tok * k + slot) * ff + row,
FC_Q4K_RBLOCK_BYTES, fc_q4k_sb_lane_grp2s)
}
/* Batched (prefill / multi-user decode) F16 gate+up — fc_moe_gateup_f16 with
* grid.z = token, exactly as _q3k_nx is to _q3k. Per token bit-exact with the
* single-token F16 kernel (same warp body). */
__global__ void fc_moe_gateup_f16_nx(const unsigned char *gate_all,
const unsigned char *up_all,
const float *x, const int *top_idx,
float *hid, int k, int ff, int cols,
int row_stride) {
extern __shared__ float sh_x[];
FC_MOE_GATEUP_NX_PROLOGUE
FC_MOE_GATEUP_BODY(xt, ((size_t)tok * k + slot) * ff + row,
FC_F16_RBLOCK_BYTES, fc_f16_sb_lane_grp2s)
}
/* Batched BF16 gate+up — fc_moe_gateup_bf16 with grid.z = token. */
__global__ void fc_moe_gateup_bf16_nx(const unsigned char *gate_all,
const unsigned char *up_all,
const float *x, const int *top_idx,
float *hid, int k, int ff, int cols,
int row_stride) {
extern __shared__ float sh_x[];
FC_MOE_GATEUP_NX_PROLOGUE
FC_MOE_GATEUP_BODY(xt, ((size_t)tok * k + slot) * ff + row,
FC_BF16_RBLOCK_BYTES, fc_bf16_sb_lane_grp2s)
}
/* Batched uint2 gate+up. These exist so the vec switch is coherent across the
* whole MoE block: with only the single-token pair vectorized, a batched run
* would silently fall back to the scalar lane map and be compared against the
* wrong replica. Per token bit-exact with fc_moe_gateup_{f16,bf16}_vec. */
__global__ void fc_moe_gateup_f16_vec_nx(const unsigned char *gate_all,
const unsigned char *up_all,
const float *x, const int *top_idx,
float *hid, int k, int ff, int cols,
int row_stride) {
extern __shared__ float4 sh_x4[];
float *sh_x = (float *)sh_x4;
FC_MOE_GATEUP_NX_PROLOGUE
FC_MOE_GATEUP_BODY(xt, ((size_t)tok * k + slot) * ff + row,
FC_F16_RBLOCK_BYTES, fc_f16_sb_lane_grp2vs)
}
__global__ void fc_moe_gateup_bf16_vec_nx(const unsigned char *gate_all,
const unsigned char *up_all,
const float *x, const int *top_idx,
float *hid, int k, int ff, int cols,
int row_stride) {
extern __shared__ float4 sh_x4[];
float *sh_x = (float *)sh_x4;
FC_MOE_GATEUP_NX_PROLOGUE
FC_MOE_GATEUP_BODY(xt, ((size_t)tok * k + slot) * ff + row,
FC_BF16_RBLOCK_BYTES, fc_bf16_sb_lane_grp2vs)
}
/*
* Batched (prefill) grouped down + weighted accumulate: grid.y = token t,
* the FC_MOE_DOWN_BODY otherwise unchanged on token t's slices —
* hid [n][k*ff], top_idx/top_wt [n][k], out [n][rows]. Per token bit-exact.
*/
#define FC_MOE_DOWN_NX_PROLOGUE \
int tok = blockIdx.y; \
hid += (size_t)tok * k * ff; \
top_idx += (size_t)tok * k; \
top_wt += (size_t)tok * k; \
out += (size_t)tok * rows;
__global__ void fc_moe_down_q3k_nx(const unsigned char *down_all,
const float *hid, const int *top_idx,
const float *top_wt, float *out, int k,
int rows, int ff, int row_stride, int wpp) {
FC_MOE_DOWN_NX_PROLOGUE
FC_MOE_DOWN_BODY(FC_Q3K_RBLOCK_BYTES, fc_q3k_sb_lane_grp2)
}
__global__ void fc_moe_down_q4k_nx(const unsigned char *down_all,
const float *hid, const int *top_idx,
const float *top_wt, float *out, int k,
int rows, int ff, int row_stride, int wpp) {
FC_MOE_DOWN_NX_PROLOGUE
FC_MOE_DOWN_BODY(FC_Q4K_RBLOCK_BYTES, fc_q4k_sb_lane_grp2)
}
__global__ void fc_moe_down_q6k_nx(const unsigned char *down_all,
const float *hid, const int *top_idx,
const float *top_wt, float *out, int k,
int rows, int ff, int row_stride, int wpp) {
FC_MOE_DOWN_NX_PROLOGUE
FC_MOE_DOWN_BODY(FC_Q6K_RBLOCK_BYTES, fc_q6k_sb_lane_grp2)
}
__global__ void fc_moe_down_f16_nx(const unsigned char *down_all,
const float *hid, const int *top_idx,
const float *top_wt, float *out, int k,
int rows, int ff, int row_stride, int wpp) {
FC_MOE_DOWN_NX_PROLOGUE
FC_MOE_DOWN_BODY(FC_F16_RBLOCK_BYTES, fc_f16_sb_lane_grp2)
}
__global__ void fc_moe_down_bf16_nx(const unsigned char *down_all,
const float *hid, const int *top_idx,
const float *top_wt, float *out, int k,
int rows, int ff, int row_stride,
int wpp) {
FC_MOE_DOWN_NX_PROLOGUE
FC_MOE_DOWN_BODY(FC_BF16_RBLOCK_BYTES, fc_bf16_sb_lane_grp2)
}
__global__ void fc_moe_down_f16_vec_nx(const unsigned char *down_all,
const float *hid, const int *top_idx,
const float *top_wt, float *out, int k,
int rows, int ff, int row_stride,
int wpp) {
FC_MOE_DOWN_NX_PROLOGUE
FC_MOE_DOWN_BODY(FC_F16_RBLOCK_BYTES, fc_f16_sb_lane_grp2v)
}
__global__ void fc_moe_down_bf16_vec_nx(const unsigned char *down_all,
const float *hid, const int *top_idx,
const float *top_wt, float *out, int k,
int rows, int ff, int row_stride,
int wpp) {
FC_MOE_DOWN_NX_PROLOGUE
FC_MOE_DOWN_BODY(FC_BF16_RBLOCK_BYTES, fc_bf16_sb_lane_grp2v)
}
/*
* Fused SwiGLU elementwise: out[i] = silu(a[i]) * b[i]. In-place safe
* (out may alias a or b). Replaces the fc_silu + fc_mul pair for the
* shared expert. silu matches fc_silu exactly: v / (1 + expf(-v)).
*
* Launch: grid ceil(n / blockDim.x), block 256 threads, no shared memory.
*/
__global__ void fc_moe_silu_mul(const float *a, const float *b, float *out,
int n) {
int i = blockIdx.x * blockDim.x + threadIdx.x;
if (i < n) {
float v = a[i];
out[i] = (v / (1.0f + expf(-v))) * b[i];
}
}
/*
* Shared-expert scalar gate, fully on-device:
* gate_logit = gate_w . x (tree-reduced dot over n)
* shexp[i] *= sigmoid(gate_logit)
* Replaces the per-layer host sync + D2H of x + host dot + fc_moe_shared_gate
* launch. gate_w is the F32 [n_embd] ffn_gate_inp_shexp vector.
*
* Launch: 1 block, 256 threads, shared 256 * sizeof(float).
*/
/* Batched (prefill) shared-expert scalar gate: one token per BLOCK
* (blockIdx.x = t) on x/shexp [n_tok][n]. Same blockDim as the single
* launch -> same reduction tree -> bit-exact per token. */
__global__ void fc_moe_shexp_dot_gate_nx(const float *gate_w, const float *x,
float *shexp, int n) {
extern __shared__ float sdata[];
int tok = blockIdx.x;
const float *xt = x + (size_t)tok * n;
float *st = shexp + (size_t)tok * n;
int tid = threadIdx.x;
int threads = blockDim.x;
float p = 0.0f;
for (int i = tid; i < n; i += threads)
p += gate_w[i] * xt[i];
sdata[tid] = p;
__syncthreads();
for (int s = threads / 2; s > 0; s >>= 1) {
if (tid < s)
sdata[tid] += sdata[tid + s];
__syncthreads();
}
float sg = 1.0f / (1.0f + expf(-sdata[0]));
for (int i = tid; i < n; i += threads)
st[i] *= sg;
}
__global__ void fc_moe_shexp_dot_gate(const float *gate_w, const float *x,
float *shexp, int n) {
extern __shared__ float sdata[];
int tid = threadIdx.x;
int threads = blockDim.x;
float p = 0.0f;
for (int i = tid; i < n; i += threads)
p += gate_w[i] * x[i];
sdata[tid] = p;
__syncthreads();
for (int s = threads / 2; s > 0; s >>= 1) {
if (tid < s)
sdata[tid] += sdata[tid + s];
__syncthreads();
}
float sgate = 1.0f / (1.0f + expf(-sdata[0]));
for (int i = tid; i < n; i += threads)
shexp[i] *= sgate;
}
/* ============================================================
* MoE expert-union dedup for the batched gate+up
* ============================================================
*
* Q3_K ONLY. Every kernel in this section decodes Q3_K (fc_q3k_decode_sb) and
* there is deliberately no Q4_K twin: dedup is default-off (QWEN35_MOE_DEDUP=1)
* and prefill-only, so it is an opt-in speed lever rather than a format the
* engine must support. The host gates on the gate/up tensor being Q3_K and
* falls back to fc_moe_gateup_q4k_nx otherwise (qwen35moe_forward_cuda.c,
* gpu_moe_batch_use_dedup) — it must never reach these with Q4_K bytes.
*
* A batched step of n positions selects k experts each — n*k routed-expert
* applications. The plain _nx gateup kernel reads (and Q3_K-decodes) every
* (token, slot) expert row separately, so an expert shared by several
* positions is re-read and re-decoded once per selecting position.
*
* Dedup collapses that to ONE weight read+decode per UNIQUE expert per step:
* the union of routed experts across the batch is materialised as per-expert
* request lists, then the grouped kernel streams each unique expert's rows
* once and applies them to every token that routed to it (grouped-by-expert
* GEMM). A token selects a given expert at most once (its k slots are
* distinct), so an expert's request count R is bounded by ntok.
*
* The saving is n*k / U, where U = |union of routed experts|. That ratio is
* what decides whether this pays, and it is set entirely by ntok:
*
* - PREFILL chunk, ntok=128, k=8, n_expert=256: 1024 applications over
* U ~ 251 experts -> ~4x fewer decodes. This is the intended caller.
* - Multi-user DECODE, ntok=8: 64 applications over U ~ 57 -> ~1.1x, which
* does not cover the wider grid.y (n_expert blocks instead of k*ntok).
* Measured as a 7-9% regression; the Pantheon caller keeps decode on _nx.
*
* Bit-exact with fc_moe_gateup_q3k_nx: the per-request hid value reuses the
* identical per-element FMA order of fc_q3k_sb_lane_grp2s — the Q3_K field
* products fs*fq are just hoisted out of the request loop (independent of x,
* so the value is unchanged), and the per-super-block d-scaled accumulate is
* reproduced verbatim per request. x is read from GLOBAL via __ldg here (vs
* the nx kernel's shared staging); the loaded float value — and thus every
* op — is identical, so hid is bit-for-bit the nx result.
*/
/* Request-tile depth: per-request accumulators held in registers. grid.z
* tiles an expert's request list by this and weights re-stream once per tile,
* so an expert with R requests is read ceil(R/RT) times.
*
* grid.z is sized off maxr (the SAFE bound on R, == ntok), not off the actual
* max R, because grid dims are host-side and the counts are device-side. At
* ntok=128 that is 16 z-tiles while the mean R is ~4, so ~15/16 of the blocks
* find r0 >= R and exit on their first load. That waste is the price of not
* syncing the counts back to the host; it is bounded by a global load plus a
* branch per block. Keep in sync with the host wrapper (fc_cuda.h
* FC_MOE_DEDUP_RT). */
#define FC_MOE_DEDUP_RT 8
/* The x-staging loop of the _shx body, in four forms; FC_MOE_SHX_STAGE selects
* one and defaults to the shipped FLAT form. Included here rather than at the
* top of the file because the forms are written in terms of FC_MOE_DEDUP_RT. */
#include "moe_dedup_shx_staging.cuh"
/*
* COMPACTED-GRID ARM, compile-time, DEFAULT OFF. At 1 it adds two kernels at the
* end of this file — fc_moe_dedup_build_wl and fc_moe_gateup_q3k_dedup_shxc —
* that replace the (n_expert x ceil(maxr/RT)) grid above with one axis over a
* device-built (expert, request-tile) work list. See the block comment there for
* the defect, the fix and the bit-exactness argument.
*
* Compile-time because the kernels ship pre-compiled in cuda/fc_cuda_ptx.h.
*
* ★ SHIPPED DEFAULT since 2026-08-11: 1, so the two kernels are IN the
* regenerated header. This is the compile-time half of a DUAL gate — the header
* only makes them available, and cuda/fc_cuda.c
* fc_moe_dedup_compact_enabled_ctx() still requires FC_MOE_DEDUP_COMPACT=1 in
* the environment to select them, which is where the safe default now lives. A
* plain build therefore behaves exactly as before unless the env asks for it.
* Gated at 1.195x in-model (recovered by Amdahl), bit-exact:
* results/prefill_moe_shx_staging_gate.md.
*
* It became a source default because the alternative was worse. Keeping this at
* 0 meant the DEPLOYED header could not be produced from a clean checkout —
* the deployed artifact needed
* make NVCCFLAGS_EXTRA="-DFC_MOE_DEDUP_COMPACT=1 -DFC_MOE_SHX_STAGE=2"
* a recipe that lived only in prose, so a plain `make` silently regressed the deployed build
* by two levers at once and neither failure was loud: this gate's runtime lookup
* is NULL-tolerant, and FC_MOE_SHX_STAGE only changes a kernel BODY. Both
* defaults now live in the source and `make` alone reproduces the shipped
* module. Set -DFC_MOE_DEDUP_COMPACT=0 for a header without the two kernels.
*
* Verify MODULE-SCOPED, never header-wide: extract FC_PTX_MOE_EXPERT_GROUP and
* md5 it. Shipped (compact+VEC4) = cb2e87d100eed7f49fe5eaeb5586fa9a, 14 .entry.
* See prefill_speed_spec.md §3.
*/
#ifndef FC_MOE_DEDUP_COMPACT
#define FC_MOE_DEDUP_COMPACT 1
#endif
/*
* Build per-expert request lists from the batch's top_idx.
* top_idx: [ntok*k] int32, flat request f = tok*k + slot -> expert id
* ecount: [n_expert] int32 (output) — tokens routing to each expert
* ereq: [n_expert*maxr] int32 (output) — ereq[e*maxr+j] = j-th request f
* maxr caps the per-expert list (== ntok). Request order within an expert is
* atomicAdd race order — order-free for the gateup consumer (each request
* writes its own hid slot). Single block; ntok*k and n_expert are both tiny.
*/
__global__ void fc_moe_dedup_build(const int *top_idx, int *ecount, int *ereq,
int n_expert, int k, int ntok, int maxr) {
int tid = threadIdx.x;
int threads = blockDim.x;
for (int e = tid; e < n_expert; e += threads)
ecount[e] = 0;
__syncthreads();
int total = ntok * k;
for (int f = tid; f < total; f += threads) {
int e = top_idx[f];
if (e < 0 || e >= n_expert)
continue;
int slot = atomicAdd(&ecount[e], 1);
if (slot < maxr)
ereq[(size_t)e * maxr + slot] = f;
}
}
/* Decode one Q3_K super-block's 8 field products fs*fq (per lane), returning
* the block scale d. Reproduces the exact per-i products of the grp2s lane
* dot; hoisting them out of the request loop is x-independent so the values
* (and the FMA that consumes them) are unchanged. */
__device__ __forceinline__ float fc_q3k_decode_sb(const unsigned char *blk,
int lane, float *wprod) {
unsigned int wv = ((const unsigned int *)blk)[lane];
float d = *(const float *)(blk + 144);
uint2 sw = ((const uint2 *)(blk + 128))[lane >> 4];
unsigned int wh = wv >> 16;
#pragma unroll
for (int i = 0; i < 8; i++) {
int qp = (i <= 4) ? 4 * i : 4 * i - 16;
unsigned int q = (i <= 4) ? wv : wh;
unsigned int s = (i < 4) ? sw.x : sw.y;
float fq = FC_MAGIC_NIB(q, qp, 4);
float fs = __fsub_rn(FC_CVT_BYTE(s, i & 3), 32.0f);
wprod[i] = __fmul_rn(fs, fq);
}
return d;
}
/*
* Grouped gate+up SwiGLU over the UNION of routed experts (dedup path).
*
* gate_all/up_all: [n_expert*ff*row_stride] repacked Q3_K (merged 3D)
* x: [ntok*cols] fp32 activations (per-token x_norm)
* ecount/ereq: per-expert request lists from fc_moe_dedup_build
* hid: [ntok*k*ff] fp32 output, hid[f*ff + row] (f = tok*k+slot)
* cols: n_embd (2048, multiple of 256); row_stride = cols/256*152
*
* Grid: (ceil(ff/8), n_expert, ceil(maxr/RT)); block 256 (8 warps), shared 0.
* blockIdx.y = expert e; warp w handles FF row blockIdx.x*8+w. Experts with
* no requests (ecount==0) exit immediately, so the full n_expert grid.y just
* skips the ~n_expert-U empty experts cheaply. For each super-block the two
* expert rows are decoded ONCE (fc_q3k_decode_sb) then applied to every
* request in the tile — one weight stream, many tokens.
*/
__global__ void fc_moe_gateup_q3k_dedup(const unsigned char *gate_all,
const unsigned char *up_all,
const float *x, const int *ecount,
const int *ereq, float *hid, int k,
int ff, int cols, int row_stride,
int maxr) {
int warp = threadIdx.x >> 5;
int lane = threadIdx.x & 31;
int row = blockIdx.x * 8 + warp;
int e = blockIdx.y;
int R = ecount[e];
int r0 = blockIdx.z * FC_MOE_DEDUP_RT;
if (r0 >= R)
return; /* no requests for e in this tile */
int rc = R - r0;
if (rc > FC_MOE_DEDUP_RT)
rc = FC_MOE_DEDUP_RT; /* uniform across the block */
int active = row < ff;
int fj[FC_MOE_DEDUP_RT];
int tokj[FC_MOE_DEDUP_RT];
float pg[FC_MOE_DEDUP_RT], pu[FC_MOE_DEDUP_RT];
#pragma unroll
for (int j = 0; j < FC_MOE_DEDUP_RT; j++) {
int f = (j < rc) ? ereq[(size_t)e * maxr + r0 + j] : -1;
fj[j] = f;
tokj[j] = (f >= 0) ? f / k : 0;
pg[j] = 0.0f;
pu[j] = 0.0f;
}
size_t expert_stride = (size_t)ff * (size_t)row_stride;
const unsigned char *gb = gate_all + (size_t)e * expert_stride +
(size_t)row * (size_t)row_stride;
const unsigned char *ub = up_all + (size_t)e * expert_stride +
(size_t)row * (size_t)row_stride;
if (active) {
int nsb = cols >> 8;
for (int sb = 0; sb < nsb; sb++) {
float wg[8], wu[8];
float dg = fc_q3k_decode_sb(gb, lane, wg);
float du = fc_q3k_decode_sb(ub, lane, wu);
#pragma unroll
for (int j = 0; j < FC_MOE_DEDUP_RT; j++) {
if (j < rc) {
const float *xb =
x + (size_t)tokj[j] * cols + (size_t)sb * 256;
float ag = 0.0f, au = 0.0f;
#pragma unroll
for (int i = 0; i < 8; i++) {
float xv = __ldg(xb + 32 * i + lane);
ag = __fmaf_rn(wg[i], xv, ag);
au = __fmaf_rn(wu[i], xv, au);
}
pg[j] += dg * ag;
pu[j] += du * au;
}
}
gb += FC_Q3K_RBLOCK_BYTES;
ub += FC_Q3K_RBLOCK_BYTES;
}
}
#pragma unroll
for (int j = 0; j < FC_MOE_DEDUP_RT; j++) {
float g = fc_warp_sum_grp(pg[j]);
float u = fc_warp_sum_grp(pu[j]);
if (active && lane == 0 && j < rc)
hid[(size_t)fj[j] * ff + row] = (g / (1.0f + expf(-g))) * u;
}
}
/*
* fc_moe_gateup_q3k_dedup_shx — the dedup gate+up with x staged in SHARED,
* the way fc_moe_gateup_q3k_nx already does it. Bit-exact with both.
*
* Two measured problems with the kernel above, fixed together because one
* change addresses both (see results/moe_dedup_occupancy_probe.md):
*
* - x traffic. All 8 warps of a block need the SAME rc rows of x; they
* differ only in which FF row they own. Reading x with __ldg from global
* therefore issues the loads 8x over, and each one carries a TEXDEPBAR.
* Counted in SASS, the per-request apply costs 51.5 instructions against
* _nx's 26 — the gap is entirely this. Staging one super-block slice per
* request (RT*256 floats = 8 KB) and reading LDS collapses it.
* - registers. Holding fj/tokj per thread costs 16 registers for values that
* are BLOCK-UNIFORM (they depend only on e, r0 and j). Moving them to 64
* bytes of shared drops the kernel from 95 registers to 62, which is
* 8 blocks/SMX instead of 5 — full occupancy.
*
* Measured on a K80 at ntok=128, k=8, n_expert=256: +20% on the kernel
* against fc_moe_gateup_q3k_dedup, bit-exact. Off by default; opt in with
* FC_MOE_DEDUP_SHX=1.
*
* Same grid, same args, same semantics as fc_moe_gateup_q3k_dedup; it just
* needs 8 KB + 64 B of static shared.
*
* The BODY is a macro so the compacted-grid entry point below can reuse it
* rather than carry a second copy. A macro and not an inline function on
* purpose: it expands to the same token stream, so the regenerated PTX for THIS
* kernel stays byte-identical to the checked-in cuda/fc_cuda_ptx.h — that
* identity is the gate on the refactor. An inline function taking the three
* __shared__ arrays as pointers would risk generic-address lowering and a
* different register allocation.
*/
/*
* The dedup gate+up work for ONE (expert, request-tile) work item.
*
* Expands inside a kernel with fc_moe_gateup_q3k_dedup_shx's signature, and
* requires these names in scope: gate_all, up_all, x, ereq, hid, k, ff, cols,
* row_stride, maxr (parameters); warp, lane, row, active (block-local); and the
* three shared arrays sh_x, sh_fj, sh_tok. E/R0/RC are the work item — expert,
* first request index, request count in the tile (1..FC_MOE_DEDUP_RT).
*
* RC gates only `j < rc`; it never appears in an arithmetic expression, and the
* per-request accumulation over super-blocks is independent of the other
* requests in the tile. So WHICH tile a request lands in cannot change its
* value — which is what makes the compacted grid below bit-exact.
*
* PRE and STAGE are the x-staging form (moe_dedup_shx_staging.cuh): PRE() emits
* whatever must be hoisted out of the super-block loop, STAGE(RC, sb) the
* staging body itself. FC_MOE_DEDUP_SHX_BODY passes the pair FC_MOE_SHX_STAGE
* selects, so the shipped kernels track the flag; the microbenchmark expands
* _BODY_S directly, once per form, to get every arm from ONE source in ONE
* translation unit.
*/
#define FC_MOE_DEDUP_SHX_BODY_S(PRE, STAGE, E, R0, RC) \
do { \
if ((int)threadIdx.x < FC_MOE_DEDUP_RT) { \
int j = threadIdx.x; \
int f = (j < (RC)) ? ereq[(size_t)(E) * maxr + (R0) + j] : -1; \
sh_fj[j] = f; \
sh_tok[j] = (f >= 0) ? f / k : 0; \
} \
__syncthreads(); \
PRE() \
\
float pg[FC_MOE_DEDUP_RT], pu[FC_MOE_DEDUP_RT]; \
_Pragma("unroll") \
for (int j = 0; j < FC_MOE_DEDUP_RT; j++) { \
pg[j] = 0.0f; \
pu[j] = 0.0f; \
} \
\
size_t expert_stride = (size_t)ff * (size_t)row_stride; \
const unsigned char *gb = gate_all + (size_t)(E) * expert_stride + \
(size_t)row * (size_t)row_stride; \
const unsigned char *ub = up_all + (size_t)(E) * expert_stride + \
(size_t)row * (size_t)row_stride; \
\
int nsb = cols >> 8; \
for (int sb = 0; sb < nsb; sb++) { \
__syncthreads(); \
STAGE(RC, sb) \
__syncthreads(); \
if (active) { \
float wg[8], wu[8]; \
float dg = fc_q3k_decode_sb(gb, lane, wg); \
float du = fc_q3k_decode_sb(ub, lane, wu); \
_Pragma("unroll") \
for (int j = 0; j < FC_MOE_DEDUP_RT; j++) { \
if (j < (RC)) { \
const float *xb = sh_x + j * 256; \
float ag = 0.0f, au = 0.0f; \
_Pragma("unroll") \
for (int i = 0; i < 8; i++) { \
float xv = xb[32 * i + lane]; \
ag = __fmaf_rn(wg[i], xv, ag); \
au = __fmaf_rn(wu[i], xv, au); \
} \
pg[j] += dg * ag; \
pu[j] += du * au; \
} \
} \
gb += FC_Q3K_RBLOCK_BYTES; \
ub += FC_Q3K_RBLOCK_BYTES; \
} \
} \
\
_Pragma("unroll") \
for (int j = 0; j < FC_MOE_DEDUP_RT; j++) { \
float g = fc_warp_sum_grp(pg[j]); \
float u = fc_warp_sum_grp(pu[j]); \
if (active && lane == 0 && j < (RC)) \
hid[(size_t)sh_fj[j] * ff + row] = \
(g / (1.0f + expf(-g))) * u; \
} \
} while (0)
/* What the shipped kernels expand: the staging form FC_MOE_SHX_STAGE selects. */
#define FC_MOE_DEDUP_SHX_BODY(E, R0, RC) \
FC_MOE_DEDUP_SHX_BODY_S(FC_MOE_SHX_STAGE_PRE, FC_MOE_SHX_STAGE_X, \
E, R0, RC)
__global__ void fc_moe_gateup_q3k_dedup_shx(const unsigned char *gate_all,
const unsigned char *up_all,
const float *x, const int *ecount,
const int *ereq, float *hid, int k,
int ff, int cols, int row_stride,
int maxr) {
__shared__ FC_MOE_SHX_ALIGN float sh_x[FC_MOE_DEDUP_RT * 256];
__shared__ int sh_fj[FC_MOE_DEDUP_RT];
__shared__ int sh_tok[FC_MOE_DEDUP_RT];
int warp = threadIdx.x >> 5;
int lane = threadIdx.x & 31;
int row = blockIdx.x * 8 + warp;
int e = blockIdx.y;
int R = ecount[e];
int r0 = blockIdx.z * FC_MOE_DEDUP_RT;
if (r0 >= R)
return; /* no requests for e in this tile */
int rc = R - r0;
if (rc > FC_MOE_DEDUP_RT)
rc = FC_MOE_DEDUP_RT; /* uniform across the block */
int active = row < ff;
FC_MOE_DEDUP_SHX_BODY(e, r0, rc);
}
#if FC_MOE_DEDUP_COMPACT
/*
* ===========================================================================
* COMPACTED GRID — compiled ONLY into a
* cd cuda_kernels && make NVCCFLAGS_EXTRA=-DFC_MOE_DEDUP_COMPACT=1
* build, so the default fc_cuda_ptx.h is byte-identical without it, and then
* opted into at runtime with FC_MOE_DEDUP_COMPACT=1. Two gates, both default
* to the safe state; with the default PTX the runtime flag finds a NULL
* CUfunction and the wrapper keeps the uncompacted launch.
* ===========================================================================
*
* THE DEFECT. fc_moe_gateup_q3k_dedup{,_shx} launch
* (ceil(ff/8), n_expert, ceil(maxr/RT)) = (64, 256, 16) = 262,144 blocks, of
* which only ceil(ff/8) * sum_e ceil(R_e/RT) do any work — 16,448, i.e. 6.3 %,
* at the shipped prefill shape. The rest load one int, find no work, exit.
* MEASURED (results/moe_dedup_occupancy_probe.md §5.2): launched with every
* ecount zero, so that NO block does any arithmetic at all, the kernel still
* costs 18.6 % of its full runtime. grid.y * grid.z is 4096 rather than ~257
* because grid dims are host-side while the per-expert counts are device-side.
*
* THE FIX. Fold (expert, request-tile) onto ONE grid axis over a device-built
* work list, and launch grid.y at the exact host-computable upper bound on the
* number of work items (fc_moe_dedup_wmax, ../cuda/fc_cuda.h): 352 instead of
* 4096, i.e. 22,528 blocks instead of 262,144 — 11.6x fewer, and 73 % of them
* useful instead of 6.3 %. NO device->host sync is introduced: the bound is
* closed-form in (n_expert, k, ntok, maxr).
*
* WHY IT IS BIT-EXACT — the whole argument:
* - The SET of (expert, r0, rc) work items is identical. The uncompacted grid
* runs (e, z) for every e with z*RT < ecount[e]; the work list holds exactly
* those, with the same r0 = z*RT and rc = min(ecount[e] - r0, RT).
* - The body is the SAME TOKENS (FC_MOE_DEDUP_SHX_BODY), over the same ereq,
* x and weights, so a work item produces the same floats either way.
* - Every output hid[f*ff + row] is written by exactly ONE work item, with a
* plain store. A request f is in exactly one expert's list exactly once, so
* it falls in exactly one tile. No atomic, no accumulation, no
* read-modify-write across blocks: the order blocks run in is not
* observable. That is stronger than "same accumulation order" — there is no
* shared accumulation to order.
* - rc is a loop bound, never an operand.
* The work list is built in expert order by an exclusive scan, so it is also
* deterministic run to run — which the shipped ereq is not (atomicAdd order),
* and does not need to be, for the same reason.
*/
/* Expert count the single-block scan below can hold. Beyond it
* fc_moe_dedup_wmax returns 0 and the caller keeps the uncompacted launch. */
#define FC_MOE_WL_EXPERT_CAP 1024
/* One (expert, request-tile) work item, packed so a gate+up block learns
* everything from ONE dependent global load — the same single load the
* uncompacted kernel spends on ecount[e]:
* .x = expert, or -1 for "past the end of the list" (the launch window is the
* host's upper bound, so its tail is sentinel-filled)
* .y = r0 << 16 | rc (r0 < maxr <= 65535, 1 <= rc <= FC_MOE_DEDUP_RT)
* Keep the maxr limit in sync with fc_moe_dedup_wmax in ../cuda/fc_cuda.c. */
#define FC_MOE_WL_PACK(r0, rc) (((r0) << 16) | (rc))
#define FC_MOE_WL_R0(y) ((y) >> 16)
#define FC_MOE_WL_RC(y) ((y) & 0xFFFF)
/*
* Compact the per-expert request counts into a (expert, request-tile) work list.
* ecount: [n_expert] int32, from fc_moe_dedup_build (which this does NOT
* modify — it runs as a separate launch so that kernel's PTX is
* untouched)
* wlist: [wmax] int2 (output) — the work items, then a (-1, 0) sentinel tail
* Single block, like fc_moe_dedup_build: n_expert is small and this is ~0.05 %
* of the gate+up launch it feeds.
*
* Thread 0 scans serially over shared memory. n_expert is at most
* FC_MOE_WL_EXPERT_CAP, and the alternative — a shuffle or tree scan — buys
* microseconds on a launch this small while re-importing a runtime-bounded
* reduction tree, the exact defect class this tree has spent a campaign
* deleting.
*/
__global__ void fc_moe_dedup_build_wl(const int *ecount, int2 *wlist,
int n_expert, int wmax) {
__shared__ int sh_wbase[FC_MOE_WL_EXPERT_CAP + 1];
int tid = threadIdx.x;
int threads = blockDim.x;
if (n_expert > FC_MOE_WL_EXPERT_CAP)
return; /* the host guarantees this; be explicit */
/* Sentinel the whole launch window FIRST, then barrier: the real entries
* below overwrite a prefix of it and must not race the fill. */
for (int w = tid; w < wmax; w += threads)
wlist[w] = make_int2(-1, 0);
for (int e = tid; e < n_expert; e += threads) {
int r = ecount[e];
sh_wbase[e] = (r > 0) ? (r + FC_MOE_DEDUP_RT - 1) / FC_MOE_DEDUP_RT : 0;
}
__syncthreads();
if (tid == 0) {
int acc = 0;
for (int e = 0; e < n_expert; e++) {
int t = sh_wbase[e];
sh_wbase[e] = acc;
acc += t;
}
sh_wbase[n_expert] = acc; /* == the real work-item count */
}
__syncthreads();
for (int e = tid; e < n_expert; e += threads) {
int r = ecount[e];
int w = sh_wbase[e];
for (int r0 = 0; r0 < r; r0 += FC_MOE_DEDUP_RT, w++) {
int rc = r - r0;
if (rc > FC_MOE_DEDUP_RT)
rc = FC_MOE_DEDUP_RT;
if (w < wmax) /* wmax >= sum ceil(R/RT) by construction */
wlist[w] = make_int2(e, FC_MOE_WL_PACK(r0, rc));
}
}
}
/*
* fc_moe_gateup_q3k_dedup_shx over the COMPACTED work list.
*
* Grid: (ceil(ff/8), wmax, 1); block 256; 8 KB + 64 B static shared — the same
* shape, block and shared as _shx with grid.y*grid.z collapsed onto grid.y.
* ecount is not read: rc travels in the work item, so a block still pays
* exactly ONE dependent load before it can tell whether it has work.
*/
__global__ void fc_moe_gateup_q3k_dedup_shxc(const unsigned char *gate_all,
const unsigned char *up_all,
const float *x, const int *ereq,
const int2 *wlist, float *hid,
int k, int ff, int cols,
int row_stride, int maxr) {
__shared__ FC_MOE_SHX_ALIGN float sh_x[FC_MOE_DEDUP_RT * 256];
__shared__ int sh_fj[FC_MOE_DEDUP_RT];
__shared__ int sh_tok[FC_MOE_DEDUP_RT];
int warp = threadIdx.x >> 5;
int lane = threadIdx.x & 31;
int row = blockIdx.x * 8 + warp;
int2 item = wlist[blockIdx.y];
int e = item.x;
if (e < 0)
return; /* past the end of the work list */
int r0 = FC_MOE_WL_R0(item.y);
int rc = FC_MOE_WL_RC(item.y);
int active = row < ff;
FC_MOE_DEDUP_SHX_BODY(e, r0, rc);
}
#endif /* FC_MOE_DEDUP_COMPACT */
} /* extern "C" */
Reference in New Issue View Git Blame Copy Permalink