8e6adf20e6
Major changes based on upstream_ref analysis: 1. factor_moe_topk_softmax.cu v2.0: Rewritten using ds_vllm/TRT-LLM warp shuffle pattern (from topk_softmax_kernels.cu). Key differences: - Zero shared memory (all butterfly __shfl_xor_sync) - VPT=2, THREADS_PER_ROW=32 (1 warp per token row) - 4 warps per CTA (4 tokens per block) - Iterative argmax with winner suppression for top-K - NaN/Inf clamping to 0 (prevents duplicate expert IDs) 2. GDN: FlashQLA backend (PROVEN on real BI-V100): - Compiles with corex clang/16 --cuda-gpu-arch=ivcore10 - Real test: NaN=False on gdn_forward(B=1, T=64, H=4, K=128) - Replaces custom factor_gdn_chunk_fwd.cu (archived to .ref) - patch_model.py now JIT-loads FlashQLA extension at runtime 3. build.sh: Correct corex flags from real compile log: --cuda-gpu-arch=ivcore10 (NOT sm_70) -D__ILUVATAR__ -D__ILUVATAR_WORKAROUND__ -D__ILUVATAR_DIAG__ -cl-single-precision-constant -mllvm --bonus-inst-threshold=0 Key insight from xllm/kernels/ilu/ixformer.h: ixformer::infer::topk_softmax() EXISTS at C++ level but Python ixformer.functions binding is missing. Our .so factor bypasses the missing Python binding entirely via dlopen/ctypes.
283 lines
10 KiB
Plaintext
283 lines
10 KiB
Plaintext
// ex_engine/csrc/factor_gdn_chunk_fwd.cu
|
||
//
|
||
// Factor 5: GDN_CHUNK_FWD — GatedDeltaNet chunked prefill forward
|
||
//
|
||
// CCCL reference: cub/device/dispatch/tuning/tuning_scan.cuh
|
||
// ScanLookbackPolicy with decoupled lookback for streaming prefix ops.
|
||
// GDN is fundamentally a recurrent scan: state[t] = decay * state[t-1] + write
|
||
//
|
||
// The NaN problem (from dockerrizhi.txt):
|
||
// "NaN in prefill GatedDeltaNet layer 0 (frac=0.9998), replacing with zeros"
|
||
// Root cause: _torch_chunk_gated_delta_rule does cumsum on gate values
|
||
// that can overflow float16 range. The FlashQLA SM70 kernel compiled but
|
||
// also produced NaN because it uses float16 accumulators.
|
||
//
|
||
// Fix: Full float32 accumulation in the recurrent state update.
|
||
// state = beta * (k ⊗ v) + exp(gate) * state [all in fp32]
|
||
// output = (q @ state).to(fp16) [cast only at output]
|
||
//
|
||
// BI-V100 tuning (SM70, 16 SMs):
|
||
// chunk_size = 16 (reduced from 64 to prevent overflow)
|
||
// head_dim = 128
|
||
// num_heads = 2 per TP rank (8 total / 4 TP)
|
||
// SMEM: state matrix = 128×128×4 = 64KB → won't fit in 48KB SMEM
|
||
// Solution: Tile state update, keep running state in registers/global
|
||
|
||
#include <cuda_runtime.h>
|
||
#include <cuda_fp16.h>
|
||
#include <float.h>
|
||
#include <math.h>
|
||
#include <stdint.h>
|
||
|
||
extern "C" {
|
||
#include "ex_engine.h"
|
||
}
|
||
|
||
// ---------------------------------------------------------------------------
|
||
// GDN Recurrent state update kernel (one CTA per head)
|
||
//
|
||
// For each chunk of tokens:
|
||
// For each time step t in chunk:
|
||
// decay = exp(gate[t]) — scalar per head
|
||
// beta_t = sigmoid(beta[t]) — scalar per head
|
||
// k_t = key[t] — (D,) vector
|
||
// v_t = value[t] — (D,) vector
|
||
// state = decay * state + beta_t * outer(k_t, v_t) — (D, D) matrix
|
||
// output[t] = query[t] @ state — (D,) vector
|
||
//
|
||
// State matrix is D×D = 128×128 = 16K floats = 64KB in fp32.
|
||
// Cannot fit in SMEM (48KB). Use register tiling: each thread owns
|
||
// a (D/TILE) × (D/TILE) block of the state matrix.
|
||
// ---------------------------------------------------------------------------
|
||
|
||
static constexpr int HEAD_DIM = 128;
|
||
static constexpr int CHUNK_SIZE = 16;
|
||
|
||
// Tile config: 256 threads, each owns a 8×8 block of state
|
||
// 128/8 = 16 tiles per dim → 16×16 = 256 tiles = 256 threads ✓
|
||
static constexpr int TILE = 8;
|
||
static constexpr int TILES_PER_DIM = HEAD_DIM / TILE; // 16
|
||
static constexpr int BLOCK_THREADS = TILES_PER_DIM * TILES_PER_DIM; // 256
|
||
|
||
__global__ void gdn_chunk_fwd_kernel(
|
||
half* __restrict__ output, // (B, L, H, D)
|
||
float* __restrict__ state_out, // (B, H, D, D) — updated state
|
||
const half* __restrict__ query, // (B, L, H, D)
|
||
const half* __restrict__ key, // (B, L, H, D)
|
||
const half* __restrict__ value, // (B, L, H, D)
|
||
const float* __restrict__ gate, // (B, L, H)
|
||
const float* __restrict__ beta, // (B, L, H)
|
||
const float* __restrict__ state_in, // (B, H, D, D) — initial state
|
||
int B, int L, int H, int D
|
||
) {
|
||
// Block: (batch, head) pair
|
||
int bh = blockIdx.x;
|
||
int b = bh / H;
|
||
int h = bh % H;
|
||
if (b >= B) return;
|
||
|
||
int tid = threadIdx.x;
|
||
int tile_row = tid / TILES_PER_DIM; // which row tile (0..15)
|
||
int tile_col = tid % TILES_PER_DIM; // which col tile (0..15)
|
||
|
||
// Each thread owns TILE×TILE = 8×8 = 64 floats of state
|
||
float my_state[TILE][TILE];
|
||
|
||
// Load initial state
|
||
int row_start = tile_row * TILE;
|
||
int col_start = tile_col * TILE;
|
||
const float* sin = state_in + (b * H + h) * D * D;
|
||
#pragma unroll
|
||
for (int r = 0; r < TILE; r++) {
|
||
#pragma unroll
|
||
for (int c = 0; c < TILE; c++) {
|
||
my_state[r][c] = sin[(row_start + r) * D + (col_start + c)];
|
||
}
|
||
}
|
||
|
||
// Shared memory for broadcast: one time step at a time
|
||
__shared__ float s_k[HEAD_DIM]; // current key vector
|
||
__shared__ float s_v[HEAD_DIM]; // current value vector
|
||
__shared__ float s_decay; // exp(gate)
|
||
__shared__ float s_beta; // sigmoid(beta)
|
||
|
||
// Process each time step sequentially (recurrent)
|
||
for (int t = 0; t < L; t++) {
|
||
// Thread 0 loads gate, beta; all threads load their k/v slice
|
||
if (tid == 0) {
|
||
float g = gate[(b * L + t) * H + h];
|
||
float bt = beta[(b * L + t) * H + h];
|
||
// Clamp gate to prevent overflow: exp(88) ≈ FLT_MAX for float32
|
||
g = fminf(fmaxf(g, -20.0f), 20.0f);
|
||
s_decay = expf(g);
|
||
s_beta = 1.0f / (1.0f + expf(-bt)); // sigmoid
|
||
}
|
||
|
||
// Cooperatively load k and v vectors into SMEM
|
||
if (tid < D) {
|
||
int idx = ((b * L + t) * H + h) * D + tid;
|
||
s_k[tid] = __half2float(key[idx]);
|
||
s_v[tid] = __half2float(value[idx]);
|
||
}
|
||
__syncthreads();
|
||
|
||
float decay = s_decay;
|
||
float bt = s_beta;
|
||
|
||
// State update: state = decay * state + beta * outer(k, v)
|
||
// Each thread updates its TILE×TILE block
|
||
#pragma unroll
|
||
for (int r = 0; r < TILE; r++) {
|
||
float k_r = s_k[row_start + r];
|
||
#pragma unroll
|
||
for (int c = 0; c < TILE; c++) {
|
||
float v_c = s_v[col_start + c];
|
||
my_state[r][c] = decay * my_state[r][c] + bt * k_r * v_c;
|
||
}
|
||
}
|
||
|
||
// Query @ state → output[t]
|
||
// Each thread computes partial dot product for its tile rows
|
||
// output[d] = sum_j query[j] * state[d][j]
|
||
// Thread (tile_row, tile_col) has state[row_start..+TILE][col_start..+TILE]
|
||
// It contributes: for each r in 0..TILE-1:
|
||
// partial[row_start+r] += sum_{c=0..TILE-1} query[col_start+c] * state[r][c]
|
||
|
||
// Load query
|
||
__shared__ float s_q[HEAD_DIM];
|
||
if (tid < D) {
|
||
int idx = ((b * L + t) * H + h) * D + tid;
|
||
s_q[tid] = __half2float(query[idx]);
|
||
}
|
||
__syncthreads();
|
||
|
||
// Compute partial result for my tile rows
|
||
float partial[TILE];
|
||
#pragma unroll
|
||
for (int r = 0; r < TILE; r++) {
|
||
partial[r] = 0.0f;
|
||
#pragma unroll
|
||
for (int c = 0; c < TILE; c++) {
|
||
partial[r] += s_q[col_start + c] * my_state[r][c];
|
||
}
|
||
}
|
||
|
||
// Reduce across col tiles (threads with same tile_row, different tile_col)
|
||
// Use shared memory: each thread writes its partial, then tile_col=0 sums
|
||
__shared__ float s_partials[TILES_PER_DIM][TILES_PER_DIM][TILE];
|
||
// s_partials[tile_row][tile_col][r]
|
||
#pragma unroll
|
||
for (int r = 0; r < TILE; r++) {
|
||
s_partials[tile_row][tile_col][r] = partial[r];
|
||
}
|
||
__syncthreads();
|
||
|
||
// tile_col == 0 aggregates across all col tiles
|
||
if (tile_col == 0) {
|
||
float result[TILE];
|
||
#pragma unroll
|
||
for (int r = 0; r < TILE; r++) {
|
||
result[r] = 0.0f;
|
||
#pragma unroll
|
||
for (int tc = 0; tc < TILES_PER_DIM; tc++) {
|
||
result[r] += s_partials[tile_row][tc][r];
|
||
}
|
||
}
|
||
// Write output
|
||
int out_base = ((b * L + t) * H + h) * D + row_start;
|
||
#pragma unroll
|
||
for (int r = 0; r < TILE; r++) {
|
||
output[out_base + r] = __float2half(result[r]);
|
||
}
|
||
}
|
||
__syncthreads();
|
||
}
|
||
|
||
// Write final state
|
||
float* sout = state_out + (b * H + h) * D * D;
|
||
#pragma unroll
|
||
for (int r = 0; r < TILE; r++) {
|
||
#pragma unroll
|
||
for (int c = 0; c < TILE; c++) {
|
||
sout[(row_start + r) * D + (col_start + c)] = my_state[r][c];
|
||
}
|
||
}
|
||
}
|
||
|
||
// ---------------------------------------------------------------------------
|
||
// Factor dispatch
|
||
// ---------------------------------------------------------------------------
|
||
|
||
static int gdn_chunk_fwd_dispatch(
|
||
void* output,
|
||
const void* input,
|
||
const void* aux_inputs[],
|
||
int n_aux,
|
||
const int64_t dims[],
|
||
int n_dims,
|
||
void* stream
|
||
) {
|
||
// dims = {B, L, H, D}
|
||
// input = query (B, L, H, D) half
|
||
// aux[0] = key, aux[1] = value, aux[2] = gate (float), aux[3] = beta (float)
|
||
// aux[4] = state_in (B, H, D, D) float
|
||
// aux[5] = state_out (B, H, D, D) float (output)
|
||
if (n_dims < 4 || n_aux < 6) return -1;
|
||
|
||
int B = (int)dims[0];
|
||
int L = (int)dims[1];
|
||
int H = (int)dims[2];
|
||
int D = (int)dims[3];
|
||
|
||
if (D != HEAD_DIM) return -1; // Only support D=128
|
||
|
||
half* out = (half*)output;
|
||
const half* q = (const half*)input;
|
||
const half* k = (const half*)aux_inputs[0];
|
||
const half* v = (const half*)aux_inputs[1];
|
||
const float* g = (const float*)aux_inputs[2];
|
||
const float* bt = (const float*)aux_inputs[3];
|
||
const float* si = (const float*)aux_inputs[4];
|
||
float* so = (float*)aux_inputs[5];
|
||
|
||
cudaStream_t cu_stream = (cudaStream_t)stream;
|
||
|
||
// Dynamic SMEM: s_partials needs TILES_PER_DIM × TILES_PER_DIM × TILE × sizeof(float)
|
||
// = 16 × 16 × 8 × 4 = 8192 bytes
|
||
// + s_k, s_v, s_q = 3 × 128 × 4 = 1536 bytes
|
||
// + s_decay, s_beta = 8 bytes
|
||
// Total ≈ 9736 bytes << 48KB ✓
|
||
|
||
dim3 grid(B * H);
|
||
dim3 block(BLOCK_THREADS); // 256
|
||
|
||
gdn_chunk_fwd_kernel<<<grid, block, 0, cu_stream>>>(
|
||
out, so, q, k, v, g, bt, si, B, L, H, D
|
||
);
|
||
|
||
return 0;
|
||
}
|
||
|
||
// ---------------------------------------------------------------------------
|
||
// .so export
|
||
// ---------------------------------------------------------------------------
|
||
|
||
static ex_factor_t s_factor;
|
||
|
||
extern "C" ex_factor_t* ex_get_factor(const ex_hardware_t* hw) {
|
||
s_factor.factor_id = EX_FACTOR_GDN_CHUNK_FWD;
|
||
s_factor.name = "gdn_chunk_fwd";
|
||
s_factor.version = "1.0.0";
|
||
s_factor.tuning = (ex_tuning_t){
|
||
.threads_per_block = BLOCK_THREADS, // 256
|
||
.items_per_thread = TILE * TILE, // 64 (state elements per thread)
|
||
.vec_size = 1,
|
||
.shared_mem_bytes = 10240, // ~10KB
|
||
.num_warps = 8,
|
||
.num_stages = 1 // sequential recurrence, no pipelining
|
||
};
|
||
s_factor.kernel = gdn_chunk_fwd_dispatch;
|
||
s_factor.kernel_fallback = NULL;
|
||
return &s_factor;
|
||
}
|