wehub-resource-sync
997 lines
43 KiB
Python
997 lines
43 KiB
Python
"""CUDA candidate for chunkwise BiGDN forward (RTX 5090 / sm_120).
|
|
|
|
Staged port. `run_cuda(inp, mode)` lets us swap each phase Triton<->CUDA:
|
|
mode="c" : Triton A + Triton B + CUDA Phase-C-fused(+divide) [Stage 1]
|
|
mode="bc" : Triton A + CUDA B + CUDA C [Stage 2]
|
|
mode="all" : full CUDA [Stage 3]
|
|
Correctness is validated phase-by-phase against Triton in the harness.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
|
|
import torch
|
|
from torch.utils.cpp_extension import load_inline
|
|
|
|
from diffusion.model.ops.fused_gdn_chunkwise import phase_a, phase_b_triton, phase_c
|
|
|
|
_EXT = None
|
|
|
|
CUDA_SRC = r"""
|
|
#include <torch/extension.h>
|
|
#include <ATen/cuda/CUDAContext.h>
|
|
#include <cuda_bf16.h>
|
|
#include <cuda_runtime.h>
|
|
#include <cuda_pipeline.h>
|
|
#include <mma.h>
|
|
|
|
using namespace nvcuda;
|
|
using bf16 = __nv_bfloat16;
|
|
|
|
namespace {
|
|
|
|
constexpr int D = 128;
|
|
constexpr int WM = 16; // wmma tile
|
|
constexpr int KT = D / WM; // 8 k-tiles
|
|
constexpr int NCT = D / WM; // 8 column tiles (output dim)
|
|
constexpr int ROWS = 64; // s-rows per CTA
|
|
constexpr int RT = ROWS / WM; // 4 row tiles
|
|
constexpr int WARPS = 8;
|
|
constexpr int THREADS = WARPS * 32;
|
|
|
|
// ───────────────────────── Phase C + fused divide ─────────────────────────
|
|
// Grid (BH*F, S_TILES): one CTA per (bh,f,s-tile of ROWS rows). M kept OUT of
|
|
// smem (it is bf16 in HBM, ~14MB total -> L2-resident); wmma b-fragments load
|
|
// directly from global so occupancy is set by the small Q working set, not by
|
|
// the 32KB M. den[r] = sum_d Q*z; out = num/(den+eps) bf16 to [B,N,H,D].
|
|
__global__ void phase_c_fused_kernel(
|
|
const bf16* __restrict__ qkv, // [B,N,3,H,D]
|
|
long sb, long sn, long s3, long sh, long sd,
|
|
const float* __restrict__ q_inv_rms, // [B,N]
|
|
const float* __restrict__ q_norm_w, // [H*D]
|
|
const float* __restrict__ rope_cos, // [N,D]
|
|
const float* __restrict__ rope_sin, // [N,D]
|
|
const bf16* __restrict__ M, // [BH,F,D,D] bf16
|
|
const float* __restrict__ z, // [BH,F,D]
|
|
bf16* __restrict__ out, // [B,N,H,D]
|
|
int B, int H, int F, int S, int S_TILES, float eps) {
|
|
const int bhf = blockIdx.x; // 0..BH*F-1
|
|
const int s_tile = blockIdx.y; // 0..S_TILES-1
|
|
const int bh = bhf / F;
|
|
const int f = bhf % F;
|
|
const int b = bh / H;
|
|
const int h = bh % H;
|
|
const int N = F * S;
|
|
const int tid = threadIdx.x;
|
|
const int warp = tid >> 5;
|
|
const int lane = tid & 31;
|
|
const int s0 = s_tile * ROWS;
|
|
|
|
extern __shared__ char smem_raw[];
|
|
float* z_s = reinterpret_cast<float*>(smem_raw); // D floats
|
|
float* qnw_s = z_s + D; // D floats
|
|
bf16* Qrot_s = reinterpret_cast<bf16*>(qnw_s + D); // ROWS*D bf16
|
|
float* den_s = reinterpret_cast<float*>(Qrot_s + ROWS * D); // ROWS floats
|
|
|
|
const bf16* M_f = M + (long)bhf * D * D;
|
|
const float* z_f = z + (long)bhf * D;
|
|
|
|
for (int i = tid; i < D; i += THREADS) { z_s[i] = z_f[i]; qnw_s[i] = q_norm_w[h * D + i]; }
|
|
__syncthreads();
|
|
|
|
const bf16* qkv_q = qkv + (long)b * sb + 0 * s3 + (long)h * sh;
|
|
const int n_base = f * S;
|
|
|
|
// ── Qrot = relu(Qn)*cos + relu(Qn[d^1])*sin, all in fp32 (match Triton) ──
|
|
for (int idx = tid; idx < ROWS * D; idx += THREADS) {
|
|
const int r = idx >> 7, d = idx & 127, s = s0 + r;
|
|
float qrot = 0.f;
|
|
if (s < S) {
|
|
const int n = n_base + s;
|
|
const float invrms = q_inv_rms[b * N + n];
|
|
const long base = (long)n * sn;
|
|
float qd = __bfloat162float(qkv_q[base + (long)d * sd]) * invrms * qnw_s[d];
|
|
float qd1 = __bfloat162float(qkv_q[base + (long)(d ^ 1) * sd]) * invrms * qnw_s[d ^ 1];
|
|
qd = qd > 0.f ? qd : 0.f;
|
|
qd1 = qd1 > 0.f ? qd1 : 0.f;
|
|
qrot = qd * rope_cos[(long)n * D + d] + qd1 * rope_sin[(long)n * D + d];
|
|
}
|
|
Qrot_s[idx] = __float2bfloat16(qrot);
|
|
}
|
|
// den uses fp32 Q (match Triton) — recompute qv from global (qkv is L1/L2-hot).
|
|
for (int r = warp; r < ROWS; r += WARPS) {
|
|
const int s = s0 + r;
|
|
float acc = 0.f;
|
|
if (s < S) {
|
|
const int n = n_base + s;
|
|
const float invrms = q_inv_rms[b * N + n];
|
|
for (int d = lane; d < D; d += 32) {
|
|
float qn = __bfloat162float(qkv_q[(long)n * sn + (long)d * sd]) * invrms * qnw_s[d];
|
|
float qv = qn > 0.f ? qn : 0.f;
|
|
acc += qv * z_s[d];
|
|
}
|
|
#pragma unroll
|
|
for (int o = 16; o > 0; o >>= 1) acc += __shfl_down_sync(0xffffffff, acc, o);
|
|
}
|
|
if (lane == 0) den_s[r] = acc;
|
|
}
|
|
__syncthreads();
|
|
|
|
// ── num = Qrot @ M (b-frag from global/L2). warp w owns col-tile w. ──
|
|
const int ct = warp; // 0..7
|
|
__shared__ float tile_smem[WARPS][WM * WM];
|
|
wmma::fragment<wmma::accumulator, WM, WM, WM, float> acc[RT];
|
|
#pragma unroll
|
|
for (int rt = 0; rt < RT; ++rt) wmma::fill_fragment(acc[rt], 0.f);
|
|
#pragma unroll
|
|
for (int kt = 0; kt < KT; ++kt) {
|
|
wmma::fragment<wmma::matrix_b, WM, WM, WM, bf16, wmma::row_major> b_frag;
|
|
wmma::load_matrix_sync(b_frag, M_f + (kt * WM) * D + ct * WM, D);
|
|
#pragma unroll
|
|
for (int rt = 0; rt < RT; ++rt) {
|
|
wmma::fragment<wmma::matrix_a, WM, WM, WM, bf16, wmma::row_major> a_frag;
|
|
wmma::load_matrix_sync(a_frag, Qrot_s + (rt * WM) * D + kt * WM, D);
|
|
wmma::mma_sync(acc[rt], a_frag, b_frag, acc[rt]);
|
|
}
|
|
}
|
|
#pragma unroll
|
|
for (int rt = 0; rt < RT; ++rt) {
|
|
wmma::store_matrix_sync(tile_smem[warp], acc[rt], WM, wmma::mem_row_major);
|
|
__syncwarp();
|
|
for (int i = lane; i < WM * WM; i += 32) {
|
|
const int rr = i >> 4, cc = i & 15;
|
|
const int s = s0 + rt * WM + rr;
|
|
if (s < S) {
|
|
const int n = n_base + s;
|
|
const int d = ct * WM + cc;
|
|
out[(((long)b * N + n) * H + h) * D + d] =
|
|
__float2bfloat16(tile_smem[warp][i] / (den_s[rt * WM + rr] + eps));
|
|
}
|
|
}
|
|
__syncwarp();
|
|
}
|
|
}
|
|
|
|
// ───────────────────────── Phase A — KV stream ─────────────────────────
|
|
// Per (bh,f): I_P_kv[D,D]=I-K_rot^T diag(b) K_rot, A[D,D]=K_rot^T diag(b) V.
|
|
// grid (BH*F,), 16 warps (512 thr). warp w owns (output=w>>3, row-tile=w&7) and
|
|
// its 8 col-tiles -> 8 acc frags/warp (avoids register spills). Krot prep shared.
|
|
constexpr int AW_KV = 16; // warp w owns one output's row-tile (8 frags)
|
|
constexpr int ATH_KV = AW_KV * 32; // 512 threads
|
|
constexpr int ACHUNK = 32; // S-rows staged per sync cycle
|
|
constexpr int ASUB = ACHUNK / WM; // wmma k-subtiles per chunk
|
|
__global__ __launch_bounds__(ATH_KV) void phase_a_kv_kernel(
|
|
const bf16* __restrict__ qkv, long sb, long sn, long s3, long sh, long sd,
|
|
const float* __restrict__ beta, // [BH, F*S]
|
|
const float* __restrict__ k_inv_rms, // [B, N]
|
|
const float* __restrict__ k_norm_w, // [H*D]
|
|
const float* __restrict__ rope_cos, // [N,D]
|
|
const float* __restrict__ rope_sin, // [N,D]
|
|
bf16* __restrict__ I_P_kv, // [BH,F,D,D]
|
|
bf16* __restrict__ A, // [BH,F,D,D]
|
|
int B, int H, int F, int S, float k_scale) {
|
|
const int bhf = blockIdx.x, bh = bhf / F, f = bhf % F, b = bh / H, h = bh % H;
|
|
const int N = F * S, tid = threadIdx.x, warp = tid >> 5, lane = tid & 31;
|
|
const int out_sel = warp >> 3; // 0=P_kv, 1=A
|
|
const int rt = warp & 7; // row-tile (d in [16*rt, 16*rt+16))
|
|
|
|
// Assumes contiguous packed qkv: sd==1, so K/V rows are contiguous in d
|
|
// (16-byte cp.async friendly). sn is the per-token stride.
|
|
extern __shared__ char sm[];
|
|
bf16* bufK = reinterpret_cast<bf16*>(sm); // [2][ACHUNK,D] raw K
|
|
bf16* bufV = bufK + 2 * ACHUNK * D; // [2][ACHUNK,D] raw V
|
|
float* knw = reinterpret_cast<float*>(bufV + 2 * ACHUNK * D); // [D]
|
|
bf16* Krot = reinterpret_cast<bf16*>(knw + D); // [ACHUNK,D]
|
|
bf16* bKrot= Krot + ACHUNK * D; // [ACHUNK,D]
|
|
bf16* bV = bKrot + ACHUNK * D; // [ACHUNK,D]
|
|
float* irms_s = reinterpret_cast<float*>(bV + ACHUNK * D); // [ACHUNK]
|
|
float* beta_s = irms_s + ACHUNK; // [ACHUNK]
|
|
|
|
for (int i = tid; i < D; i += ATH_KV) knw[i] = k_norm_w[h * D + i];
|
|
|
|
wmma::fragment<wmma::accumulator, WM, WM, WM, float> acc[NCT];
|
|
#pragma unroll
|
|
for (int ct = 0; ct < NCT; ++ct) wmma::fill_fragment(acc[ct], 0.f);
|
|
|
|
const bf16* qkv_k = qkv + (long)b * sb + 1 * s3 + (long)h * sh;
|
|
const bf16* qkv_v = qkv + (long)b * sb + 2 * s3 + (long)h * sh;
|
|
const float* beta_bhf = beta + (long)bh * N + (long)f * S;
|
|
const int n_base = f * S;
|
|
const int NCH = (S + ACHUNK - 1) / ACHUNK;
|
|
const int VEC = 8; // 8 bf16 = 16 bytes per cp.async
|
|
const int COPIES = (ACHUNK * D) / VEC; // 16-byte copies per chunk per tensor
|
|
|
|
// cp.async stage raw K,V for chunk c into double-buffer slot c%2.
|
|
#define STAGE(c) { \
|
|
const int _s0 = (c) * ACHUNK; const int _slot = ((c) & 1) * ACHUNK * D; \
|
|
for (int t = tid; t < COPIES; t += ATH_KV) { \
|
|
const int _sl = t / (D / VEC), _d = (t % (D / VEC)) * VEC, _s = _s0 + _sl; \
|
|
if (_s < S) { \
|
|
const long _ko = (long)(n_base + _s) * sn + _d; \
|
|
__pipeline_memcpy_async(&bufK[_slot + _sl * D + _d], &qkv_k[_ko], 16); \
|
|
__pipeline_memcpy_async(&bufV[_slot + _sl * D + _d], &qkv_v[_ko], 16); \
|
|
} \
|
|
} \
|
|
__pipeline_commit(); \
|
|
}
|
|
|
|
STAGE(0);
|
|
__syncthreads();
|
|
for (int c = 0; c < NCH; ++c) {
|
|
if (c + 1 < NCH) STAGE(c + 1);
|
|
__pipeline_wait_prior(c + 1 < NCH ? 1 : 0);
|
|
const int s0 = c * ACHUNK;
|
|
// hoist per-row constants (invrms, beta) once -> avoid 128x redundant loads.
|
|
for (int i = tid; i < ACHUNK; i += ATH_KV) {
|
|
const int s = s0 + i;
|
|
irms_s[i] = (s < S) ? k_inv_rms[b * N + n_base + s] : 0.f;
|
|
beta_s[i] = (s < S) ? beta_bhf[s] : 0.f;
|
|
}
|
|
__syncthreads();
|
|
const bf16* Kb = bufK + (c & 1) * ACHUNK * D;
|
|
const bf16* Vb = bufV + (c & 1) * ACHUNK * D;
|
|
// vectorized prep: each thread does a 4-wide d-block (d4..d4+3). RoPE pair
|
|
// d^1 stays within the block, so one float4 cos/sin/knw load covers all 4.
|
|
for (int t = tid; t < ACHUNK * (D / 4); t += ATH_KV) {
|
|
const int sl = t / (D / 4), d4 = (t % (D / 4)) * 4, s = s0 + sl;
|
|
float kr[4] = {0, 0, 0, 0}, bvv[4] = {0, 0, 0, 0};
|
|
float be = 0.f;
|
|
if (s < S) {
|
|
const int n = n_base + s;
|
|
const float invrms = irms_s[sl];
|
|
be = beta_s[sl];
|
|
const float4 c4 = *reinterpret_cast<const float4*>(&rope_cos[(long)n * D + d4]);
|
|
const float4 s4 = *reinterpret_cast<const float4*>(&rope_sin[(long)n * D + d4]);
|
|
const float4 w4 = *reinterpret_cast<const float4*>(&knw[d4]);
|
|
const float wv[4] = {w4.x, w4.y, w4.z, w4.w};
|
|
const float cv[4] = {c4.x, c4.y, c4.z, c4.w};
|
|
const float sv[4] = {s4.x, s4.y, s4.z, s4.w};
|
|
float kn[4];
|
|
#pragma unroll
|
|
for (int j = 0; j < 4; ++j) {
|
|
float v = __bfloat162float(Kb[sl * D + d4 + j]) * invrms * wv[j];
|
|
kn[j] = (v > 0.f ? v : 0.f) * k_scale;
|
|
}
|
|
#pragma unroll
|
|
for (int j = 0; j < 4; ++j) {
|
|
kr[j] = kn[j] * cv[j] + kn[j ^ 1] * sv[j]; // d^1 within block
|
|
bvv[j] = be * __bfloat162float(Vb[sl * D + d4 + j]);
|
|
}
|
|
}
|
|
#pragma unroll
|
|
for (int j = 0; j < 4; ++j) {
|
|
Krot[sl * D + d4 + j] = __float2bfloat16(kr[j]);
|
|
bKrot[sl * D + d4 + j] = __float2bfloat16(be * kr[j]);
|
|
bV[sl * D + d4 + j] = __float2bfloat16(bvv[j]);
|
|
}
|
|
}
|
|
__syncthreads();
|
|
const bf16* bsrc = out_sel ? bV : bKrot;
|
|
#pragma unroll
|
|
for (int ks = 0; ks < ASUB; ++ks) {
|
|
wmma::fragment<wmma::matrix_a, WM, WM, WM, bf16, wmma::col_major> a_frag;
|
|
wmma::load_matrix_sync(a_frag, Krot + ks * WM * D + rt * WM, D);
|
|
#pragma unroll
|
|
for (int ct = 0; ct < NCT; ++ct) {
|
|
wmma::fragment<wmma::matrix_b, WM, WM, WM, bf16, wmma::row_major> bfr;
|
|
wmma::load_matrix_sync(bfr, bsrc + ks * WM * D + ct * WM, D);
|
|
wmma::mma_sync(acc[ct], a_frag, bfr, acc[ct]);
|
|
}
|
|
}
|
|
__syncthreads();
|
|
}
|
|
#undef STAGE
|
|
|
|
bf16* dst = (out_sel ? A : I_P_kv) + (long)bhf * D * D;
|
|
__shared__ float st[AW_KV][WM * WM];
|
|
#pragma unroll
|
|
for (int ct = 0; ct < NCT; ++ct) {
|
|
wmma::store_matrix_sync(st[warp], acc[ct], WM, wmma::mem_row_major);
|
|
__syncwarp();
|
|
for (int i = lane; i < WM * WM; i += 32) {
|
|
const int d = rt * WM + (i >> 4), dp = ct * WM + (i & 15);
|
|
float v = st[warp][i];
|
|
if (out_sel == 0) v = (d == dp ? 1.f : 0.f) - v; // I - P_kv
|
|
dst[(long)d * D + dp] = __float2bfloat16(v);
|
|
}
|
|
__syncwarp();
|
|
}
|
|
}
|
|
|
|
// ───────────────────────── Phase A — Z stream ─────────────────────────
|
|
// Per (bh,f): I_P_z[D,D]=I-K^T diag(b) K (no RoPE), B[D]=sum_s b_s K_s.
|
|
__global__ void phase_a_z_kernel(
|
|
const bf16* __restrict__ qkv, long sb, long sn, long s3, long sh, long sd,
|
|
const float* __restrict__ beta,
|
|
const float* __restrict__ k_inv_rms,
|
|
const float* __restrict__ k_norm_w,
|
|
bf16* __restrict__ I_P_z, // [BH,F,D,D]
|
|
float* __restrict__ Bz, // [BH,F,D]
|
|
int B, int H, int F, int S, float k_scale) {
|
|
const int bhf = blockIdx.x, bh = bhf / F, f = bhf % F, b = bh / H, h = bh % H;
|
|
const int N = F * S, tid = threadIdx.x, warp = tid >> 5, lane = tid & 31;
|
|
|
|
extern __shared__ char sm[];
|
|
bf16* K_s = reinterpret_cast<bf16*>(sm); // [ACHUNK,D] bf16 (K)
|
|
bf16* bK = K_s + ACHUNK * D; // [ACHUNK,D] bf16 (beta*K)
|
|
float* knw = reinterpret_cast<float*>(bK + ACHUNK * D); // [D]
|
|
float* Bacc = knw + D; // [D]
|
|
|
|
for (int i = tid; i < D; i += THREADS) { knw[i] = k_norm_w[h * D + i]; Bacc[i] = 0.f; }
|
|
|
|
wmma::fragment<wmma::accumulator, WM, WM, WM, float> acc_p[NCT];
|
|
#pragma unroll
|
|
for (int ct = 0; ct < NCT; ++ct) wmma::fill_fragment(acc_p[ct], 0.f);
|
|
|
|
const bf16* qkv_k = qkv + (long)b * sb + 1 * s3 + (long)h * sh;
|
|
const float* beta_bhf = beta + (long)bh * N + (long)f * S;
|
|
const int n_base = f * S;
|
|
__syncthreads();
|
|
|
|
for (int s0 = 0; s0 < S; s0 += ACHUNK) {
|
|
for (int idx = tid; idx < ACHUNK * D; idx += THREADS) {
|
|
const int sl = idx >> 7, d = idx & 127, s = s0 + sl;
|
|
float kvv = 0.f, be = 0.f;
|
|
if (s < S) {
|
|
const int n = n_base + s;
|
|
float kn = __bfloat162float(qkv_k[(long)n * sn + (long)d * sd]) * k_inv_rms[b * N + n] * knw[d];
|
|
kvv = (kn > 0.f ? kn : 0.f) * k_scale;
|
|
be = beta_bhf[s];
|
|
}
|
|
K_s[idx] = __float2bfloat16(kvv);
|
|
const float bk = be * kvv;
|
|
bK[idx] = __float2bfloat16(bk);
|
|
if (s < S) atomicAdd(&Bacc[d], bk);
|
|
}
|
|
__syncthreads();
|
|
#pragma unroll
|
|
for (int ks = 0; ks < ASUB; ++ks) {
|
|
wmma::fragment<wmma::matrix_a, WM, WM, WM, bf16, wmma::col_major> a_frag;
|
|
wmma::load_matrix_sync(a_frag, K_s + ks * WM * D + warp * WM, D);
|
|
#pragma unroll
|
|
for (int ct = 0; ct < NCT; ++ct) {
|
|
wmma::fragment<wmma::matrix_b, WM, WM, WM, bf16, wmma::row_major> bp;
|
|
wmma::load_matrix_sync(bp, bK + ks * WM * D + ct * WM, D);
|
|
wmma::mma_sync(acc_p[ct], a_frag, bp, acc_p[ct]);
|
|
}
|
|
}
|
|
__syncthreads();
|
|
}
|
|
|
|
bf16* IPz = I_P_z + (long)bhf * D * D;
|
|
__shared__ float st[WARPS][WM * WM];
|
|
#pragma unroll
|
|
for (int ct = 0; ct < NCT; ++ct) {
|
|
wmma::store_matrix_sync(st[warp], acc_p[ct], WM, wmma::mem_row_major);
|
|
__syncwarp();
|
|
for (int i = lane; i < WM * WM; i += 32) {
|
|
const int d = warp * WM + (i >> 4), dp = ct * WM + (i & 15);
|
|
IPz[(long)d * D + dp] = __float2bfloat16((d == dp ? 1.f : 0.f) - st[warp][i]);
|
|
}
|
|
__syncwarp();
|
|
}
|
|
for (int i = tid; i < D; i += THREADS) Bz[(long)bhf * D + i] = Bacc[i];
|
|
}
|
|
|
|
// ═══════════════════ CAM path (live model) — no prep, [B,H,D,N] ═══════════════════
|
|
// cam_scan_bidi_chunkwise: identity norm/RoPE, skip_relu, skip_z, num_only,
|
|
// output transposed to [B,H,D,N]. K_rot==K==raw k. Reads q/k/v directly (no
|
|
// packing), writes fp32 out directly (no permute). 16 warps; warp w owns
|
|
// (out=w>>3 {P_kv,A}, row-tile=w&7).
|
|
__global__ __launch_bounds__(ATH_KV) void cam_phase_a_kv_kernel(
|
|
const float* __restrict__ k, // [B,H,D,N]
|
|
const float* __restrict__ v, // [B,H,D,N]
|
|
const float* __restrict__ beta, // [BH,F,S]
|
|
bf16* __restrict__ I_P_kv, bf16* __restrict__ A,
|
|
int B, int H, int F, int S) {
|
|
const int bhf = blockIdx.x, bh = bhf / F, f = bhf % F;
|
|
const int tid = threadIdx.x, warp = tid >> 5, lane = tid & 31;
|
|
const int out_sel = warp >> 3, rt = warp & 7;
|
|
const long Nd = (long)F * S; // tokens per (b,h)
|
|
const long kbase = (long)bh * D * Nd; // k/v base for this bh
|
|
|
|
extern __shared__ char sm[];
|
|
bf16* Ksm = reinterpret_cast<bf16*>(sm); // [ACHUNK,D]
|
|
bf16* bKsm = Ksm + ACHUNK * D; // [ACHUNK,D]
|
|
bf16* bVsm = bKsm + ACHUNK * D; // [ACHUNK,D]
|
|
|
|
wmma::fragment<wmma::accumulator, WM, WM, WM, float> acc[NCT];
|
|
#pragma unroll
|
|
for (int ct = 0; ct < NCT; ++ct) wmma::fill_fragment(acc[ct], 0.f);
|
|
|
|
const float* beta_bhf = beta + (long)bhf * S;
|
|
const int n_base = f * S;
|
|
|
|
for (int s0 = 0; s0 < S; s0 += ACHUNK) {
|
|
// stage K/V D-MAJOR [d][sl]: read along N (contiguous, coalesced — k/v are
|
|
// [B,H,D,N]) and write contiguously. a_frag = K is then row_major; the wmma
|
|
// transpose moves to the b-operand (col_major). Avoids the uncoalesced
|
|
// strided global loads that were the L1TEX bottleneck (ncu).
|
|
for (int idx = tid; idx < D * ACHUNK; idx += ATH_KV) {
|
|
const int d = idx / ACHUNK, sl = idx - d * ACHUNK, s = s0 + sl;
|
|
float kk = 0.f, bv = 0.f, bk = 0.f;
|
|
if (s < S) {
|
|
const long off = kbase + (long)d * Nd + (n_base + s);
|
|
const float be = beta_bhf[s];
|
|
kk = k[off]; bk = be * kk; bv = be * v[off];
|
|
}
|
|
Ksm[idx] = __float2bfloat16(kk); // Ksm[d*ACHUNK + sl]
|
|
bKsm[idx] = __float2bfloat16(bk);
|
|
bVsm[idx] = __float2bfloat16(bv);
|
|
}
|
|
__syncthreads();
|
|
const bf16* bsrc = out_sel ? bVsm : bKsm;
|
|
#pragma unroll
|
|
for (int ks = 0; ks < ASUB; ++ks) {
|
|
// a = K[rt-rows, ks-cols] row_major (d-major smem, ldm=ACHUNK)
|
|
wmma::fragment<wmma::matrix_a, WM, WM, WM, bf16, wmma::row_major> a_frag;
|
|
wmma::load_matrix_sync(a_frag, Ksm + rt * WM * ACHUNK + ks * WM, ACHUNK);
|
|
#pragma unroll
|
|
for (int ct = 0; ct < NCT; ++ct) {
|
|
// b = (betaK)^T : col_major from d-major betaK, ldm=ACHUNK
|
|
wmma::fragment<wmma::matrix_b, WM, WM, WM, bf16, wmma::col_major> bfr;
|
|
wmma::load_matrix_sync(bfr, bsrc + ct * WM * ACHUNK + ks * WM, ACHUNK);
|
|
wmma::mma_sync(acc[ct], a_frag, bfr, acc[ct]);
|
|
}
|
|
}
|
|
__syncthreads();
|
|
}
|
|
|
|
bf16* dst = (out_sel ? A : I_P_kv) + (long)bhf * D * D;
|
|
__shared__ float st[AW_KV][WM * WM];
|
|
#pragma unroll
|
|
for (int ct = 0; ct < NCT; ++ct) {
|
|
wmma::store_matrix_sync(st[warp], acc[ct], WM, wmma::mem_row_major);
|
|
__syncwarp();
|
|
for (int i = lane; i < WM * WM; i += 32) {
|
|
const int d = rt * WM + (i >> 4), dp = ct * WM + (i & 15);
|
|
float val = st[warp][i];
|
|
if (out_sel == 0) val = (d == dp ? 1.f : 0.f) - val;
|
|
dst[(long)d * D + dp] = __float2bfloat16(val);
|
|
}
|
|
__syncwarp();
|
|
}
|
|
}
|
|
|
|
// num = Q @ M_hist, write transposed fp32 to [B,H,D,N]. grid (BH*F, S_TILES).
|
|
__global__ void cam_phase_c_kernel(
|
|
const float* __restrict__ q, // [B,H,D,N]
|
|
const bf16* __restrict__ M, // [BH,F,D,D] bf16
|
|
float* __restrict__ out, // [B,H,D,N]
|
|
int B, int H, int F, int S, int S_TILES) {
|
|
const int bhf = blockIdx.x, s_tile = blockIdx.y, f = bhf % F;
|
|
const int bh = bhf / F;
|
|
const int tid = threadIdx.x, warp = tid >> 5, lane = tid & 31;
|
|
const int s0 = s_tile * ROWS;
|
|
const long Nd = (long)F * S;
|
|
const long qbase = (long)bh * D * Nd;
|
|
const int n_base = f * S;
|
|
|
|
extern __shared__ char smem_raw[];
|
|
bf16* Qsm = reinterpret_cast<bf16*>(smem_raw); // [D, ROWS] d-major
|
|
const bf16* M_f = M + (long)bhf * D * D;
|
|
|
|
// stage Q D-MAJOR [i][r]: read q[i, n] along N (coalesced; q is [B,H,D,N]).
|
|
// a_frag = Q is then col_major (a[s,i]=Q[i,s]); b_frag = M stays row_major.
|
|
for (int idx = tid; idx < D * ROWS; idx += THREADS) {
|
|
const int i = idx / ROWS, r = idx - i * ROWS, s = s0 + r;
|
|
float qv = 0.f;
|
|
if (s < S) qv = q[qbase + (long)i * Nd + (n_base + s)];
|
|
Qsm[idx] = __float2bfloat16(qv); // Qsm[i*ROWS + r]
|
|
}
|
|
__syncthreads();
|
|
|
|
const int ct = warp;
|
|
__shared__ float tile_smem[WARPS][WM * WM];
|
|
wmma::fragment<wmma::accumulator, WM, WM, WM, float> acc[ROWS / WM];
|
|
#pragma unroll
|
|
for (int rt = 0; rt < ROWS / WM; ++rt) wmma::fill_fragment(acc[rt], 0.f);
|
|
#pragma unroll
|
|
for (int kt = 0; kt < KT; ++kt) {
|
|
wmma::fragment<wmma::matrix_b, WM, WM, WM, bf16, wmma::row_major> b_frag;
|
|
wmma::load_matrix_sync(b_frag, M_f + (kt * WM) * D + ct * WM, D);
|
|
#pragma unroll
|
|
for (int rt = 0; rt < ROWS / WM; ++rt) {
|
|
wmma::fragment<wmma::matrix_a, WM, WM, WM, bf16, wmma::col_major> a_frag;
|
|
wmma::load_matrix_sync(a_frag, Qsm + (kt * WM) * ROWS + rt * WM, ROWS);
|
|
wmma::mma_sync(acc[rt], a_frag, b_frag, acc[rt]);
|
|
}
|
|
}
|
|
#pragma unroll
|
|
for (int rt = 0; rt < ROWS / WM; ++rt) {
|
|
// col-major store -> tile_smem[j*16 + s]; then consecutive lanes write
|
|
// consecutive n (=s) for a fixed j -> coalesced output (out is [B,H,D,N]).
|
|
wmma::store_matrix_sync(tile_smem[warp], acc[rt], WM, wmma::mem_col_major);
|
|
__syncwarp();
|
|
for (int i = lane; i < WM * WM; i += 32) {
|
|
const int s = s0 + rt * WM + (i & 15);
|
|
if (s < S) {
|
|
const int j = ct * WM + (i >> 4);
|
|
out[qbase + (long)j * Nd + (n_base + s)] = tile_smem[warp][i];
|
|
}
|
|
}
|
|
__syncwarp();
|
|
}
|
|
}
|
|
|
|
// cam Phase B: serial F scan, fwd+rev combined -> bf16 M_hist directly (no fp32
|
|
// roundtrip). skip_z. M state kept bf16 in smem. grid (BH,), 8 warps, warp w
|
|
// owns output row-tile w. M = g*(I-P_kv)@M + A.
|
|
__global__ void cam_phase_b_kernel(
|
|
const bf16* __restrict__ I_P_kv, // [BH,F,D,D]
|
|
const bf16* __restrict__ A, // [BH,F,D,D]
|
|
const float* __restrict__ decay, // [BH,F]
|
|
bf16* __restrict__ M_hist, // [BH,F,D,D]
|
|
int BH, int F) {
|
|
const int bh = blockIdx.x;
|
|
const int tid = threadIdx.x, warp = tid >> 5, lane = tid & 31;
|
|
const int rt = warp; // row-tile (8 warps -> 8 row-tiles)
|
|
extern __shared__ char smem_raw[];
|
|
bf16* Mbf = reinterpret_cast<bf16*>(smem_raw); // [D,D] running state
|
|
__shared__ float st[WARPS][WM * WM];
|
|
|
|
const bf16* IP_bh = I_P_kv + (long)bh * F * D * D;
|
|
const bf16* A_bh = A + (long)bh * F * D * D;
|
|
const float* g_bh = decay + (long)bh * F;
|
|
|
|
// helper macro replaced by inline: compute M_new for frame f into st/Mbf
|
|
#define CAM_B_FRAME(f, WRITE_ADD, dstframe) \
|
|
{ \
|
|
const bf16* IP_f = IP_bh + (long)(f) * D * D; \
|
|
const bf16* A_f = A_bh + (long)(f) * D * D; \
|
|
const float g = g_bh[(f)]; \
|
|
wmma::fragment<wmma::accumulator, WM, WM, WM, float> acc[NCT]; \
|
|
_Pragma("unroll") for (int ct = 0; ct < NCT; ++ct) wmma::fill_fragment(acc[ct], 0.f); \
|
|
_Pragma("unroll") for (int kt = 0; kt < KT; ++kt) { \
|
|
wmma::fragment<wmma::matrix_a, WM, WM, WM, bf16, wmma::row_major> a_frag; \
|
|
wmma::load_matrix_sync(a_frag, IP_f + (long)(rt * WM) * D + kt * WM, D); \
|
|
_Pragma("unroll") for (int ct = 0; ct < NCT; ++ct) { \
|
|
wmma::fragment<wmma::matrix_b, WM, WM, WM, bf16, wmma::row_major> b_frag; \
|
|
wmma::load_matrix_sync(b_frag, Mbf + (long)(kt * WM) * D + ct * WM, D); \
|
|
wmma::mma_sync(acc[ct], a_frag, b_frag, acc[ct]); \
|
|
} \
|
|
} \
|
|
__syncthreads(); /* all warps done reading Mbf before overwrite */ \
|
|
_Pragma("unroll") for (int ct = 0; ct < NCT; ++ct) { \
|
|
wmma::store_matrix_sync(st[warp], acc[ct], WM, wmma::mem_row_major); \
|
|
__syncwarp(); \
|
|
for (int i = lane; i < WM * WM; i += 32) { \
|
|
const int d = rt * WM + (i >> 4), dp = ct * WM + (i & 15); \
|
|
float mnew = g * st[warp][i] + __bfloat162float(A_f[(long)d * D + dp]); \
|
|
Mbf[(long)d * D + dp] = __float2bfloat16(mnew); \
|
|
bf16* hp = M_hist + ((long)bh * F + (dstframe)) * D * D + (long)d * D + dp; \
|
|
if (WRITE_ADD) *hp = __float2bfloat16(__bfloat162float(*hp) + mnew); \
|
|
else *hp = __float2bfloat16(mnew); \
|
|
} \
|
|
__syncwarp(); \
|
|
} \
|
|
__syncthreads(); \
|
|
}
|
|
|
|
// forward
|
|
for (int i = tid; i < D * D; i += THREADS) Mbf[i] = __float2bfloat16(0.f);
|
|
__syncthreads();
|
|
for (int f = 0; f < F; ++f) CAM_B_FRAME(f, false, f);
|
|
// reverse (combined): for f_src=F-1..1, add into M_hist[f_src-1]
|
|
for (int i = tid; i < D * D; i += THREADS) Mbf[i] = __float2bfloat16(0.f);
|
|
__syncthreads();
|
|
for (int f = F - 1; f >= 1; --f) CAM_B_FRAME(f, true, f - 1);
|
|
#undef CAM_B_FRAME
|
|
}
|
|
|
|
} // namespace
|
|
|
|
void cam_phase_b(torch::Tensor I_P_kv, torch::Tensor A, torch::Tensor decay,
|
|
torch::Tensor M_hist, int F) {
|
|
const int BH = I_P_kv.size(0);
|
|
auto stream = at::cuda::getCurrentCUDAStream();
|
|
size_t smem = sizeof(bf16) * D * D;
|
|
cam_phase_b_kernel<<<BH, THREADS, smem, stream>>>(
|
|
reinterpret_cast<const bf16*>(I_P_kv.data_ptr()),
|
|
reinterpret_cast<const bf16*>(A.data_ptr()),
|
|
decay.data_ptr<float>(), reinterpret_cast<bf16*>(M_hist.data_ptr()), BH, F);
|
|
}
|
|
|
|
void cam_phase_a_kv(torch::Tensor k, torch::Tensor v, torch::Tensor beta,
|
|
torch::Tensor I_P_kv, torch::Tensor A, int F, int S) {
|
|
const int B = k.size(0), H = k.size(1), BH = B * H;
|
|
auto stream = at::cuda::getCurrentCUDAStream();
|
|
size_t smem = sizeof(bf16) * (3 * ACHUNK * D);
|
|
cam_phase_a_kv_kernel<<<BH * F, ATH_KV, smem, stream>>>(
|
|
k.data_ptr<float>(), v.data_ptr<float>(), beta.data_ptr<float>(),
|
|
reinterpret_cast<bf16*>(I_P_kv.data_ptr()), reinterpret_cast<bf16*>(A.data_ptr()),
|
|
B, H, F, S);
|
|
}
|
|
|
|
void cam_phase_c(torch::Tensor q, torch::Tensor M, torch::Tensor out, int F, int S) {
|
|
const int B = q.size(0), H = q.size(1), BH = B * H;
|
|
const int S_TILES = (S + ROWS - 1) / ROWS;
|
|
auto stream = at::cuda::getCurrentCUDAStream();
|
|
size_t smem = sizeof(bf16) * (ROWS * D);
|
|
dim3 grid(BH * F, S_TILES);
|
|
cam_phase_c_kernel<<<grid, THREADS, smem, stream>>>(
|
|
q.data_ptr<float>(), reinterpret_cast<const bf16*>(M.data_ptr()),
|
|
out.data_ptr<float>(), B, H, F, S, S_TILES);
|
|
}
|
|
|
|
void phase_a_kv(torch::Tensor qkv, torch::Tensor beta, torch::Tensor k_inv_rms,
|
|
torch::Tensor k_norm_w, torch::Tensor rope_cos, torch::Tensor rope_sin,
|
|
torch::Tensor I_P_kv, torch::Tensor A, int F, int S, double k_scale) {
|
|
const int B = qkv.size(0), H = qkv.size(3), BH = I_P_kv.size(0);
|
|
auto stream = at::cuda::getCurrentCUDAStream();
|
|
size_t smem = sizeof(float) * (D + 2 * ACHUNK) + sizeof(bf16) * (7 * ACHUNK * D);
|
|
static bool s_kv = false;
|
|
if (!s_kv) { cudaFuncSetAttribute(phase_a_kv_kernel,
|
|
cudaFuncAttributeMaxDynamicSharedMemorySize, (int)smem); s_kv = true; }
|
|
phase_a_kv_kernel<<<BH * F, ATH_KV, smem, stream>>>(
|
|
reinterpret_cast<const bf16*>(qkv.data_ptr()),
|
|
qkv.stride(0), qkv.stride(1), qkv.stride(2), qkv.stride(3), qkv.stride(4),
|
|
beta.data_ptr<float>(), k_inv_rms.data_ptr<float>(), k_norm_w.data_ptr<float>(),
|
|
rope_cos.data_ptr<float>(), rope_sin.data_ptr<float>(),
|
|
reinterpret_cast<bf16*>(I_P_kv.data_ptr()), reinterpret_cast<bf16*>(A.data_ptr()),
|
|
B, H, F, S, (float)k_scale);
|
|
}
|
|
|
|
void phase_a_z(torch::Tensor qkv, torch::Tensor beta, torch::Tensor k_inv_rms,
|
|
torch::Tensor k_norm_w, torch::Tensor I_P_z, torch::Tensor Bz,
|
|
int F, int S, double k_scale) {
|
|
const int B = qkv.size(0), H = qkv.size(3), BH = I_P_z.size(0);
|
|
auto stream = at::cuda::getCurrentCUDAStream();
|
|
size_t smem = sizeof(bf16) * (2 * ACHUNK * D) + sizeof(float) * (2 * D);
|
|
phase_a_z_kernel<<<BH * F, THREADS, smem, stream>>>(
|
|
reinterpret_cast<const bf16*>(qkv.data_ptr()),
|
|
qkv.stride(0), qkv.stride(1), qkv.stride(2), qkv.stride(3), qkv.stride(4),
|
|
beta.data_ptr<float>(), k_inv_rms.data_ptr<float>(), k_norm_w.data_ptr<float>(),
|
|
reinterpret_cast<bf16*>(I_P_z.data_ptr()), Bz.data_ptr<float>(),
|
|
B, H, F, S, (float)k_scale);
|
|
}
|
|
|
|
void phase_c_fused(torch::Tensor qkv, torch::Tensor q_inv_rms, torch::Tensor q_norm_w,
|
|
torch::Tensor rope_cos, torch::Tensor rope_sin,
|
|
torch::Tensor M, torch::Tensor z, torch::Tensor out,
|
|
int F, int S, double eps) {
|
|
const int B = qkv.size(0);
|
|
const int H = qkv.size(3);
|
|
const int BH = M.size(0);
|
|
const int S_TILES = (S + ROWS - 1) / ROWS;
|
|
auto stream = at::cuda::getCurrentCUDAStream();
|
|
dim3 grid(BH * F, S_TILES);
|
|
size_t smem = sizeof(float) * (D + D + ROWS) + sizeof(bf16) * (ROWS * D);
|
|
phase_c_fused_kernel<<<grid, THREADS, smem, stream>>>(
|
|
reinterpret_cast<const bf16*>(qkv.data_ptr()),
|
|
qkv.stride(0), qkv.stride(1), qkv.stride(2), qkv.stride(3), qkv.stride(4),
|
|
q_inv_rms.data_ptr<float>(), q_norm_w.data_ptr<float>(),
|
|
rope_cos.data_ptr<float>(), rope_sin.data_ptr<float>(),
|
|
reinterpret_cast<const bf16*>(M.data_ptr()), z.data_ptr<float>(),
|
|
reinterpret_cast<bf16*>(out.data_ptr()),
|
|
B, H, F, S, S_TILES, (float)eps);
|
|
}
|
|
"""
|
|
|
|
CPP_SRC = r"""
|
|
#include <torch/extension.h>
|
|
void phase_c_fused(torch::Tensor qkv, torch::Tensor q_inv_rms, torch::Tensor q_norm_w,
|
|
torch::Tensor rope_cos, torch::Tensor rope_sin,
|
|
torch::Tensor M, torch::Tensor z, torch::Tensor out,
|
|
int F, int S, double eps);
|
|
void phase_a_kv(torch::Tensor qkv, torch::Tensor beta, torch::Tensor k_inv_rms,
|
|
torch::Tensor k_norm_w, torch::Tensor rope_cos, torch::Tensor rope_sin,
|
|
torch::Tensor I_P_kv, torch::Tensor A, int F, int S, double k_scale);
|
|
void phase_a_z(torch::Tensor qkv, torch::Tensor beta, torch::Tensor k_inv_rms,
|
|
torch::Tensor k_norm_w, torch::Tensor I_P_z, torch::Tensor Bz,
|
|
int F, int S, double k_scale);
|
|
void cam_phase_a_kv(torch::Tensor k, torch::Tensor v, torch::Tensor beta,
|
|
torch::Tensor I_P_kv, torch::Tensor A, int F, int S);
|
|
void cam_phase_b(torch::Tensor I_P_kv, torch::Tensor A, torch::Tensor decay,
|
|
torch::Tensor M_hist, int F);
|
|
void cam_phase_c(torch::Tensor q, torch::Tensor M, torch::Tensor out, int F, int S);
|
|
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
|
|
m.def("phase_c_fused", &phase_c_fused);
|
|
m.def("phase_a_kv", &phase_a_kv);
|
|
m.def("phase_a_z", &phase_a_z);
|
|
m.def("cam_phase_a_kv", &cam_phase_a_kv);
|
|
m.def("cam_phase_b", &cam_phase_b);
|
|
m.def("cam_phase_c", &cam_phase_c);
|
|
}
|
|
"""
|
|
|
|
|
|
def _system_cuda_home():
|
|
"""Locate a CUDA toolkit with *usable* headers (cuda_runtime.h, nv/target,
|
|
crt/host_config.h) and an nvcc, returning ``(home, include_dir)``.
|
|
|
|
Some conda nvcc installs (and the original svideo box) ship nvcc without
|
|
co-located headers; a complete toolkit may instead live under
|
|
``/usr/local/cuda-*`` (``include/`` layout) or in a conda env
|
|
(``targets/<arch>/include`` layout). We probe both, honoring an explicit
|
|
``$CUDA_HOME``/``$CONDA_PREFIX`` first so the same kernel builds on any
|
|
cluster (CW H100, GB200, RTX 5090)."""
|
|
import sys as _sys
|
|
|
|
env = [os.environ.get("CUDA_HOME"), os.environ.get("CUDA_PATH"), os.environ.get("CONDA_PREFIX"), _sys.prefix]
|
|
cands = [c for c in env if c] + ["/usr/local/cuda-12.9", "/usr/local/cuda", "/usr/local/cuda-12"]
|
|
for c in cands:
|
|
if not c or not os.path.isfile(os.path.join(c, "bin", "nvcc")):
|
|
continue
|
|
# Accept either the classic include/ layout or conda's targets/<arch>/include.
|
|
from glob import glob
|
|
|
|
inc_dirs = [os.path.join(c, "include")] + glob(os.path.join(c, "targets", "*", "include"))
|
|
for inc in inc_dirs:
|
|
if os.path.isfile(os.path.join(inc, "cuda_runtime.h")) and os.path.isfile(
|
|
os.path.join(inc, "nv", "target")
|
|
):
|
|
return c, inc
|
|
return None, None
|
|
|
|
|
|
def build(name="cw_cuda_ext_v1", extra_cuda=None):
|
|
global _EXT
|
|
if _EXT is not None:
|
|
return _EXT
|
|
# Compile for the GPU actually present (sm_90 H100, sm_100 GB200, sm_120
|
|
# 5090) so the default-on path produces a runnable cubin everywhere; an
|
|
# explicit TORCH_CUDA_ARCH_LIST still wins. The kernels are arch-portable
|
|
# (wmma + cp.async, sm_80+; no clusters/TMA).
|
|
if "TORCH_CUDA_ARCH_LIST" not in os.environ:
|
|
try:
|
|
_maj, _min = torch.cuda.get_device_capability()
|
|
os.environ["TORCH_CUDA_ARCH_LIST"] = f"{_maj}.{_min}"
|
|
except Exception:
|
|
os.environ["TORCH_CUDA_ARCH_LIST"] = "12.0"
|
|
import torch.utils.cpp_extension as C
|
|
|
|
home, inc = _system_cuda_home()
|
|
incs = []
|
|
if home:
|
|
os.environ["CUDA_HOME"] = home
|
|
os.environ["PATH"] = os.path.join(home, "bin") + os.pathsep + os.environ.get("PATH", "")
|
|
C.CUDA_HOME = home # override torch's import-time cached value
|
|
incs = [inc]
|
|
_EXT = load_inline(
|
|
name=name,
|
|
cpp_sources=CPP_SRC,
|
|
cuda_sources=CUDA_SRC,
|
|
extra_cuda_cflags=["-O3", "-lineinfo", "-Xptxas=-v"] + (extra_cuda or []),
|
|
extra_include_paths=incs,
|
|
with_cuda=True,
|
|
verbose=True,
|
|
)
|
|
return _EXT
|
|
|
|
|
|
def cuda_phase_c_fused(inp, M_hist, z_hist, eps=1e-6):
|
|
ext = build()
|
|
qkv = inp["qkv"]
|
|
B, N, _, H, Dd = qkv.shape
|
|
out = torch.empty(B, N, H, Dd, device=qkv.device, dtype=qkv.dtype)
|
|
M_bf = M_hist.to(torch.bfloat16).contiguous()
|
|
ext.phase_c_fused(
|
|
qkv.contiguous(),
|
|
inp["q_inv_rms"].contiguous(),
|
|
inp["q_norm_w"].contiguous(),
|
|
inp["rope_cos"].contiguous(),
|
|
inp["rope_sin"].contiguous(),
|
|
M_bf,
|
|
z_hist.contiguous(),
|
|
out,
|
|
inp["F"],
|
|
inp["S"],
|
|
eps,
|
|
)
|
|
return out
|
|
|
|
|
|
def cuda_phase_a(inp, k_scale=1.0):
|
|
"""CUDA Phase A -> (I_P_kv, A, I_P_z, B_z), shapes/dtypes matching Triton."""
|
|
ext = build()
|
|
qkv = inp["qkv"].contiguous()
|
|
B, N, _, H, Dd = qkv.shape
|
|
BH, F, S = B * H, inp["F"], inp["S"]
|
|
dev = qkv.device
|
|
BD = 1 << (Dd - 1).bit_length() # next pow2 (==128 for Dd<=128)
|
|
I_P_kv = torch.empty(BH, F, BD, BD, device=dev, dtype=torch.bfloat16)
|
|
A = torch.empty(BH, F, BD, BD, device=dev, dtype=torch.bfloat16)
|
|
I_P_z = torch.empty(BH, F, BD, BD, device=dev, dtype=torch.bfloat16)
|
|
B_z = torch.empty(BH, F, BD, device=dev, dtype=torch.float32)
|
|
beta = inp["beta"].contiguous().float()
|
|
kir = inp["k_inv_rms"].contiguous().float()
|
|
knw = inp["k_norm_w"].contiguous().float()
|
|
rc, rs = inp["rope_cos"].contiguous(), inp["rope_sin"].contiguous()
|
|
ext.phase_a_kv(qkv, beta, kir, knw, rc, rs, I_P_kv, A, F, S, k_scale)
|
|
ext.phase_a_z(qkv, beta, kir, knw, I_P_z, B_z, F, S, k_scale)
|
|
return I_P_kv, A, I_P_z, B_z
|
|
|
|
|
|
def cam_scan_bidi_chunkwise_cuda(q, k, v, beta, decay, *, dot_precision=None):
|
|
"""Drop-in for fused_gdn_chunkwise.cam_scan_bidi_chunkwise.
|
|
|
|
CUDA fast path when head_dim==128 & fp32 contiguous inputs (the deployed
|
|
sana_wm config: attention_head_dim=128). Falls back to the Triton entry
|
|
otherwise so this is always safe to substitute.
|
|
"""
|
|
B, H, D, N = q.shape
|
|
F = beta.shape[2]
|
|
if (
|
|
not _CUDA_CAM_DISABLED
|
|
and D == 128
|
|
and q.dtype == torch.float32
|
|
and N % F == 0
|
|
and q.is_contiguous()
|
|
and k.is_contiguous()
|
|
and v.is_contiguous()
|
|
):
|
|
try:
|
|
return run_cuda_cam(q, k, v, beta, decay)
|
|
except Exception as e: # missing toolkit / unexpected shape -> Triton
|
|
_disable_cuda_cam(e)
|
|
from diffusion.model.ops.fused_gdn_chunkwise import cam_scan_bidi_chunkwise
|
|
|
|
return cam_scan_bidi_chunkwise(q, k, v, beta, decay, dot_precision=dot_precision)
|
|
|
|
|
|
_CAM_BUFS = {}
|
|
# Triton phase_b (fp32 state) is fast (0.04ms) & accurate. The CUDA cam_phase_b
|
|
# kept bf16 M-state across frames -> drift (max_rel 0.14) and was slow (serial,
|
|
# BH CTAs) -> disabled by default. fp32-state CUDA phase B is future work.
|
|
_CAM_PHASE_B = os.environ.get("CAM_PHASE_B", "triton") # "triton" (default) or "cuda"
|
|
|
|
# Once the CUDA path fails to build/run (e.g. no nvcc / no CUDA toolkit on this
|
|
# box), latch it off so the drop-ins fall back to Triton silently for the rest
|
|
# of the process instead of retrying the (failing, slow) compile every call.
|
|
_CUDA_CAM_DISABLED = False
|
|
|
|
|
|
def _disable_cuda_cam(err):
|
|
global _CUDA_CAM_DISABLED
|
|
if not _CUDA_CAM_DISABLED:
|
|
_CUDA_CAM_DISABLED = True
|
|
print(f"[SANA_GDN_CUDA] CUDA cam kernels unavailable ({err}); using Triton")
|
|
|
|
|
|
def cam_scan_chunkwise_cuda(
|
|
q, k, v, beta, decay, *, reverse=False, init_state=None, save_final_state=False, dot_precision=None
|
|
):
|
|
"""Drop-in for fused_gdn_chunkwise.cam_scan_chunkwise (the STREAMING/cached
|
|
causal cam scan). CUDA cam_phase_a + cam_phase_c (read q/k/v [B,H,D,N] direct,
|
|
write transposed fp32 direct), Triton phase_b for the state-cached single
|
|
direction. bf16 GEMM. Falls back to Triton for unsupported shapes/dtype."""
|
|
B, H, D, N = q.shape
|
|
F = beta.shape[2]
|
|
S = N // F
|
|
use_cuda = (
|
|
not _CUDA_CAM_DISABLED
|
|
and D == 128
|
|
and q.dtype == torch.float32
|
|
and N % F == 0
|
|
and q.is_contiguous()
|
|
and k.is_contiguous()
|
|
and v.is_contiguous()
|
|
)
|
|
if use_cuda:
|
|
try:
|
|
ext = build()
|
|
BH, BD, dev = B * H, 128, q.device
|
|
I_P_kv = torch.empty(BH, F, BD, BD, device=dev, dtype=torch.bfloat16)
|
|
A = torch.empty_like(I_P_kv)
|
|
ext.cam_phase_a_kv(k.contiguous(), v.contiguous(), beta.contiguous().float(), I_P_kv, A, F, S)
|
|
dummy = torch.empty(1, device=dev, dtype=torch.float32)
|
|
init_kv = init_state.to(torch.float32).contiguous() if init_state is not None else None
|
|
init_z = torch.zeros(BH, BD, device=dev, dtype=torch.float32) if init_state is not None else None
|
|
direction = 2 if reverse else 1
|
|
pb = phase_b_triton(
|
|
I_P_kv,
|
|
A,
|
|
dummy,
|
|
dummy,
|
|
decay.contiguous(),
|
|
F=F,
|
|
dot_precision=0,
|
|
direction=direction,
|
|
init_state_kv=init_kv,
|
|
init_state_z=init_z,
|
|
return_final_state=save_final_state,
|
|
skip_z=True,
|
|
)
|
|
if save_final_state:
|
|
M_fwd, _zf, M_rev, _zr, final_kv, _fz = pb
|
|
else:
|
|
M_fwd, _zf, M_rev, _zr = pb
|
|
M_use = M_rev if reverse else M_fwd
|
|
M_bf = M_use.to(torch.bfloat16).contiguous()
|
|
out = torch.empty(B, H, D, N, device=dev, dtype=torch.float32)
|
|
ext.cam_phase_c(q.contiguous(), M_bf, out, F, S)
|
|
return (out, final_kv) if save_final_state else out
|
|
except Exception as e: # missing toolkit / unexpected shape -> Triton
|
|
_disable_cuda_cam(e)
|
|
from diffusion.model.ops.fused_gdn_chunkwise import cam_scan_chunkwise
|
|
|
|
return cam_scan_chunkwise(
|
|
q,
|
|
k,
|
|
v,
|
|
beta,
|
|
decay,
|
|
reverse=reverse,
|
|
init_state=init_state,
|
|
save_final_state=save_final_state,
|
|
dot_precision=dot_precision,
|
|
)
|
|
|
|
|
|
def run_cuda_cam(q, k, v, beta, decay, reuse_buffers=True):
|
|
"""CUDA cam path: reads q/k/v [B,H,D,N] directly, writes fp32 [B,H,D,N]
|
|
directly (no packing, no transpose). Phase B via Triton (skip_z).
|
|
|
|
reuse_buffers=True caches the I_P_kv/A/M_bf/out scratch by shape — a
|
|
streaming deployment reuses these every step, so per-call torch.empty/.to
|
|
(which otherwise hits cudaMalloc sync under memory pressure) is amortised."""
|
|
ext = build()
|
|
B, H, D, N = q.shape
|
|
BH, F = B * H, beta.shape[2]
|
|
S = N // F
|
|
dev = q.device
|
|
BD = 1 << (D - 1).bit_length()
|
|
k = k.contiguous()
|
|
v = v.contiguous()
|
|
q = q.contiguous()
|
|
beta_f = beta.contiguous().float()
|
|
key = (B, H, D, N, F, dev)
|
|
buf = _CAM_BUFS.get(key) if reuse_buffers else None
|
|
if buf is None:
|
|
buf = dict(
|
|
I_P_kv=torch.empty(BH, F, BD, BD, device=dev, dtype=torch.bfloat16),
|
|
A=torch.empty(BH, F, BD, BD, device=dev, dtype=torch.bfloat16),
|
|
M_bf=torch.empty(BH, F, BD, BD, device=dev, dtype=torch.bfloat16),
|
|
out=torch.empty(B, H, D, N, device=dev, dtype=torch.float32),
|
|
)
|
|
if reuse_buffers:
|
|
_CAM_BUFS[key] = buf
|
|
I_P_kv, A, M_bf, out = buf["I_P_kv"], buf["A"], buf["M_bf"], buf["out"]
|
|
ext.cam_phase_a_kv(k, v, beta_f, I_P_kv, A, F, S)
|
|
if _CAM_PHASE_B == "cuda":
|
|
ext.cam_phase_b(I_P_kv, A, decay.contiguous().float(), M_bf, F)
|
|
else:
|
|
dummy = torch.empty(1, device=dev, dtype=torch.float32)
|
|
M_hist, _, _, _ = phase_b_triton(
|
|
I_P_kv,
|
|
A,
|
|
dummy,
|
|
dummy,
|
|
decay.contiguous(),
|
|
F=F,
|
|
dot_precision=0,
|
|
direction=0,
|
|
combined_history=True,
|
|
skip_z=True,
|
|
)
|
|
M_bf.copy_(M_hist)
|
|
ext.cam_phase_c(q, M_bf, out, F, S)
|
|
return out
|
|
|
|
|
|
def run_cuda(inp, mode="ac", dot_precision=0, eps=1e-6):
|
|
qkv = inp["qkv"]
|
|
F, S = inp["F"], inp["S"]
|
|
if mode in ("ac", "abc_a"):
|
|
I_P_kv, A, I_P_z, B_z = cuda_phase_a(inp)
|
|
else:
|
|
I_P_kv, A, I_P_z, B_z = phase_a(
|
|
qkv,
|
|
inp["beta"],
|
|
inp["q_inv_rms"],
|
|
inp["k_inv_rms"],
|
|
inp["q_norm_w"],
|
|
inp["k_norm_w"],
|
|
inp["rope_cos"],
|
|
inp["rope_sin"],
|
|
F=F,
|
|
S=S,
|
|
dot_precision=dot_precision,
|
|
)
|
|
M_hist, z_hist, _, _ = phase_b_triton(
|
|
I_P_kv, A, I_P_z, B_z, inp["decay"], F=F, dot_precision=dot_precision, direction=0, combined_history=True
|
|
)
|
|
return cuda_phase_c_fused(inp, M_hist, z_hist, eps=eps)
|