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;
|
|||
|
|
}
|