Niko1221/Strata
1
0 Fork 0
mirror of https://github.com/Niko1221/Strata.git synced 2026-11-07 14:21:06 -10:00
Code Issues Packages Projects Releases 1 Wiki Activity

GDN TC: two value tiles per head on cards with >= 96 SMs (STRATA_GDN_TC_TILES), harness --tiles

Browse Source
This commit is contained in:
Niko1221 committed 2026-10-30 03:46:18 +01:00
1 parent 0ea3e455fd
commit 51567de1a6
4 files changed
+56 -21
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

include/strata/prefill/gdn_probe.hpp
+1
View File
@@ -30,6 +30,7 @@ struct GdnProbe {
// in
int force = kGdnAuto; // kGdnAuto: dispatch as production does; otherwise only this path may run
bool any_sms = false; // chunked: skip the 128-SM test, as STRATA_GDN_CHUNKED=2 does
int tc_tiles = 0; // tcs: value tiles a head (1, 2; 0 = STRATA_GDN_TC_TILES, else automatic: 2 on cards with >= 96 SMs)
int gsb = 0; // chunked: tokens per scratch super-block (a multiple of 32, at most 2048; 0 = 2048)
// out (set by every gdn_recurrence call made on the thread that installed the probe)
int selected = kGdnNone; // the path that launched
src/prefill/gdn_harness.cu
+3
View File
@@ -60,6 +60,7 @@ struct Args {
int reps = 7, warmup = 2;
bool any_sms = false;
int gsb = 0;
int tiles = 0;
bool fallback_incumbent = false; // a forced path takes the incumbent for calls it cannot run (mirrors production)
int64_t ld16 = 0;
FixtureSpec fx;
@@ -164,6 +165,7 @@ RunOut run_path(Dev& d, const Args& a, const std::vector<int64_t>& calls, int pa
sp::GdnProbe probe;
probe.any_sms = a.any_sms;
probe.gsb = a.gsb;
probe.tc_tiles = a.tiles;
sp::gdn_probe_install(&probe);
std::vector<cudaEvent_t> ev;
if (timed) {
@@ -415,6 +417,7 @@ int main(int argc, char** argv) {
else if (k == "--warmup") a.warmup = std::atoi(val());
else if (k == "--any-sms") a.any_sms = true;
else if (k == "--gsb") a.gsb = std::atoi(val());
else if (k == "--tiles") a.tiles = std::atoi(val());
else if (k == "--fallback-incumbent") a.fallback_incumbent = true;
else if (k == "--ld16") a.ld16 = parse_list(val())[0];
else if (k == "--state") a.fx.state = argv[++i];
src/prefill/gdn_tc.cuh
+25 -16
View File
@@ -438,7 +438,7 @@ __global__ void __launch_bounds__(256) gdn_tc_prep_kernel(const float* __restric
// 8 warps, 2 a sub-partition, may have 255).
// NORM (step 5: the output norm fused, the FP16 output written here) is not implemented yet; oc_out takes the FP32
// output as gdn_chunk_scan_kernel writes it.
template <bool NORM, bool SPLIT = false, int NP = 3>
template <bool NORM, bool SPLIT = false, int NP = 3, int NTL = 1>
__global__ void __launch_bounds__(GTC_CW * 32, 1)
gdn_tc_scan_kernel(float* __restrict__ state, const float* __restrict__ h, const uint16_t* __restrict__ tig_,
const float* __restrict__ teg, const float4* __restrict__ scg, const uint16_t* __restrict__ kqg,
@@ -449,7 +449,12 @@ gdn_tc_scan_kernel(float* __restrict__ state, const float* __restrict__ h, const
// sigmoid(z) itself (the sums of squares across the block's 8 warps through shared memory), and not the FP32 o: one pass less
// over the 24 KB-a-token output. Another summation order than the kernel's 4 partials (not bitwise).
static_assert(!NORM || SPLIT, "gdn_tc_scan_kernel: the fused output norm is implemented for the split plan");
constexpr int NT = GTC_CW * 32;
// NTL value tiles a head (a launch of HV x NTL blocks of GTC_CW / NTL warps): block (vh, tile) owns the state columns
// [tile S / NTL, (tile + 1) S / NTL), so a card with more than 48 SMs keeps more of them busy. The k / q / T / P images are staged
// by every tile of the head; v only for its columns. The fused output norm needs the whole head in one block: NORM => NTL == 1.
static_assert(!NORM || NTL == 1, "gdn_tc_scan_kernel: the fused output norm needs a whole head in a block");
static_assert(GTC_CW % NTL == 0, "gdn_tc_scan_kernel: tiles of whole warps");
constexpr int NT = (GTC_CW / NTL) * 32;
extern __shared__ __align__(16) unsigned char gtc_sm[];
uint16_t* kh = reinterpret_cast<uint16_t*>(gtc_sm); // [2][GCH][GTC_KP] f16 of k x 2^8 (hi)
uint16_t* kl = kh + 2 * GCH * GTC_KP; // [2][GCH][GTC_KP] what hi left out (lo)
@@ -464,7 +469,7 @@ gdn_tc_scan_kernel(float* __restrict__ state, const float* __restrict__ h, const
float4* sc = reinterpret_cast<float4*>(pl + NV * GCH * GTC_PP); // [2][GCH] (beta, exp(gamma), exp(gL - gamma), gamma)
float* ste = reinterpret_cast<float*>(sc + 2 * GCH); // [2] 2^-et
float* rss = ste + 2; // NORM: [GTC_CW][GCH] partial sums of squares
const int vh = blockIdx.x, qhd = vh % HK;
const int vh = blockIdx.x / NTL, qhd = vh % HK, cb = (blockIdx.x % NTL) * (S / NTL); // the tile's first column
const int tid = threadIdx.x, warp = tid >> 5, lane = tid & 31, gid = lane >> 2, tig = lane & 3;
const size_t rs = (size_t) HV * S;
// chunk ch into buffer b, this thread's 16-B pieces (one cp.async group): the four k / q images (4 x GCH rows of
@@ -479,7 +484,7 @@ gdn_tc_scan_kernel(float* __restrict__ state, const float* __restrict__ h, const
const uint16_t* img = kqg + ((size_t) ch * HK + qhd) * 4 * GCH * S;
#pragma unroll
for (int m = 0; m < 4 * GCH * S / 8 / NT; ++m) {
const int im = m >> 1, t = ((m & 1) * NT + tid) / (S / 8), r = (tid % (S / 8)) * 8;
const int pc = tid + NT * m, im = pc / (GCH * S / 8), t = (pc % (GCH * S / 8)) / (S / 8), r = (pc % (S / 8)) * 8;
uint16_t* dst = (im == 0 ? kh : im == 1 ? kl : im == 2 ? qh : ql) + (b * GCH + t) * GTC_KP + r;
gdn_cp16(reinterpret_cast<float*>(dst), reinterpret_cast<const float*>(img + (im * GCH + t) * S + r));
}
@@ -488,27 +493,31 @@ gdn_tc_scan_kernel(float* __restrict__ state, const float* __restrict__ h, const
}
if (parts & 2) {
#pragma unroll
for (int m = 0; m < GCH * S / 4 / NT; ++m) {
const int p = tid + NT * m, t = p / (S / 4), r = (p % (S / 4)) * 4;
for (int m = 0; m < GCH * (S / NTL) / 4 / NT; ++m) {
const int p = tid + NT * m, t = p / (S / NTL / 4), r = (p % (S / NTL / 4)) * 4;
float* dst = sv + (bv * GCH + t) * GTC_VP + r;
if (t < n) gdn_cp16(dst, h + (c_t0 + t) * C + 2 * HK * S + vh * S + r);
if (t < n) gdn_cp16(dst, h + (c_t0 + t) * C + 2 * HK * S + vh * S + cb + r);
else *reinterpret_cast<float4*>(dst) = make_float4(0.0f, 0.0f, 0.0f, 0.0f);
}
const int im = tid / (GCH * GCH / 8), t = (tid % (GCH * GCH / 8)) / (GCH / 8), i = (tid % (GCH / 8)) * 8;
const int64_t src = ((((int64_t) ch * HV + vh) * 2 + im) * GCH + t) * GCH + i;
const int dst = (bv * GCH + t) * GTC_PP + i;
gdn_cp16(reinterpret_cast<float*>((im ? tl : th) + dst), reinterpret_cast<const float*>(tig_ + src));
gdn_cp16(reinterpret_cast<float*>((im ? pl : ph) + dst), reinterpret_cast<const float*>(pig + src));
#pragma unroll
for (int m = 0; m < 2 * GCH * GCH / 8 / NT; ++m) {
const int pc = tid + NT * m, im = pc / (GCH * GCH / 8), t = (pc % (GCH * GCH / 8)) / (GCH / 8), i = (pc % (GCH / 8)) * 8;
const int64_t src = ((((int64_t) ch * HV + vh) * 2 + im) * GCH + t) * GCH + i;
const int dst = (bv * GCH + t) * GTC_PP + i;
gdn_cp16(reinterpret_cast<float*>((im ? tl : th) + dst), reinterpret_cast<const float*>(tig_ + src));
gdn_cp16(reinterpret_cast<float*>((im ? pl : ph) + dst), reinterpret_cast<const float*>(pig + src));
}
}
gdn_cp_commit();
};
const int c0 = 16 * warp; // the warp's value columns c0 + gid and c0 + gid + 8
const int c0 = 16 * warp; // the warp's value columns c0 + gid and c0 + gid + 8 of the tile (v in shared memory)
const int ca = cb + c0; // the same in the head (state, output)
// the state slice as S^T fragments: tile j = rows 8j..8j+7; s[j][0] = S[8j + 2tig][c0 + gid], [1] the row after,
// [2] column + 8, [3] both (64 scalar loads a lane, 32-B runs, once a launch)
float s[16][4];
#pragma unroll
for (int j = 0; j < 16; ++j) {
const float* p = state + (size_t) (8 * j + 2 * tig) * rs + (size_t) vh * S + c0 + gid;
const float* p = state + (size_t) (8 * j + 2 * tig) * rs + (size_t) vh * S + ca + gid;
s[j][0] = p[0]; s[j][1] = p[rs]; s[j][2] = p[8]; s[j][3] = p[rs + 8];
}
stage(0, 0, 3);
@@ -619,7 +628,7 @@ gdn_tc_scan_kernel(float* __restrict__ state, const float* __restrict__ h, const
#pragma unroll
for (int q = 0; q < 4; ++q) {
const int t = 8 * nt + 2 * tig + (q & 1);
if (t < n) oc_out[(c_t0 + t) * rs + (size_t) vh * S + c0 + gid + 8 * (q >> 1)] = (ot[nt][q] * sOi) * out_scale;
if (t < n) oc_out[(c_t0 + t) * rs + (size_t) vh * S + ca + gid + 8 * (q >> 1)] = (ot[nt][q] * sOi) * out_scale;
}
}
if (SPLIT) { // v, T and P of this chunk are dead: every warp is past (d); the next chunk's land during (e)
@@ -667,7 +676,7 @@ gdn_tc_scan_kernel(float* __restrict__ state, const float* __restrict__ h, const
}
#pragma unroll
for (int j = 0; j < 16; ++j) {
float* p = state + (size_t) (8 * j + 2 * tig) * rs + (size_t) vh * S + c0 + gid;
float* p = state + (size_t) (8 * j + 2 * tig) * rs + (size_t) vh * S + ca + gid;
p[0] = s[j][0]; p[rs] = s[j][1]; p[8] = s[j][2]; p[rs + 8] = s[j][3];
}
}
src/prefill/kernels.cu
+27 -5
View File
@@ -1279,6 +1279,24 @@ cudaError_t gdn_rec_chunked(float* state, const float* h, const float* gate, con
// STRATA_GDN_TC=1: gdn_rec_chunked with the scan on the tensor cores (gdn_tc_scan_kernel, gdn_tc.cuh) and the prep
// that also writes its f16 images (gdn_tc_prep_kernel, scratch gdn_tc_scratch), 48 blocks of 256 threads in one wave. An error return means nothing was launched: the card is below
// sm_80, has fewer than 48 SMs, or less than 122 KB of shared memory a block.
// value tiles a head for the split tensor-core scan: STRATA_GDN_TC_TILES=1|2 (or the probe's tc_tiles) asks; 1 otherwise
int gdn_tc_tiles() {
const GdnProbe* pr = gdn_probe_get();
if (pr != nullptr && pr->tc_tiles > 0) return pr->tc_tiles >= 2 ? 2 : 1;
static const int e = [] { const char* v = std::getenv("STRATA_GDN_TC_TILES"); return v ? std::atoi(v) : 0; }();
if (e > 0) return e >= 2 ? 2 : 1;
// automatic: two tiles when all 2 x 48 blocks fit the card in one wave (measured: PRO 5000 / 110 SMs rec 1.47x; A6000 / 84 SMs
// is 1.25x SLOWER with two tiles, its 96 blocks need a second wave)
static int sms_of[64] = {};
int dev = 0;
if (cudaGetDevice(&dev) != cudaSuccess || dev < 0 || dev >= 64) { cudaGetLastError(); return 1; }
if (sms_of[dev] == 0) {
int n = 0;
if (cudaDeviceGetAttribute(&n, cudaDevAttrMultiProcessorCount, dev) != cudaSuccess) { cudaGetLastError(); n = 1; }
sms_of[dev] = n;
}
return sms_of[dev] >= 2 * HV ? 2 : 1;
}
// the output norm's operands, when the split scan fuses it (see gdn_tc_scan_kernel NORM)
struct GdnTcNorm {
const float* z;
@@ -1288,7 +1306,7 @@ struct GdnTcNorm {
int64_t ld16;
};
cudaError_t gdn_rec_tc(float* state, const float* h, const float* gate, const float* beta, float* y, int64_t T,
cudaStream_t s, bool split = false, int np = 3, const GdnTcNorm* nm = nullptr) {
cudaStream_t s, bool split = false, int np = 3, const GdnTcNorm* nm = nullptr, int tiles = 1) {
if (T <= 0) return cudaSuccess;
{
int dev = 0, major = 0, sms = 0, smem = 0;
@@ -1313,6 +1331,7 @@ cudaError_t gdn_rec_tc(float* state, const float* h, const float* gate, const fl
int dev = 0;
cudaGetDevice(&dev);
if (dev >= 0 && dev < 64 && !attrs[ai][dev]) {
if (split) cudaFuncSetAttribute(gdn_tc_scan_kernel<false, true, 3, 2>, cudaFuncAttributeMaxDynamicSharedMemorySize, (int) kTcScanSmemSplit);
if (ai == 2) {
cudaFuncSetAttribute(gdn_tc_scan_kernel<true, true>, cudaFuncAttributeMaxDynamicSharedMemorySize, (int) kTcScanSmemSplitNorm);
cudaFuncSetAttribute(gdn_tc_scan_kernel<true, true, 2>, cudaFuncAttributeMaxDynamicSharedMemorySize, (int) kTcScanSmemSplitNorm);
@@ -1329,7 +1348,9 @@ cudaError_t gdn_rec_tc(float* state, const float* h, const float* gate, const fl
for (int64_t t0 = 0; t0 < T; t0 += GSB) {
const int nch = (int) ((std::min<int64_t>(GSB, T - t0) + GCH - 1) / GCH);
gdn_tc_prep_kernel<<<dim3((unsigned) nch, HK), 256, kTcPrepSmem, s>>>(h, gate, beta, sc.ti, sc.te, sc.sc, sc.kq, sc.pi, t0, T);
if (nm && split) {
if (tiles == 2 && split && !nm && np == 3) {
gdn_tc_scan_kernel<false, true, 3, 2><<<HV * 2, GTC_CW * 16, kTcScanSmemSplit, s>>>(state, h, sc.ti, sc.te, sc.sc, sc.kq, sc.pi, y, t0, nch, T);
} else if (nm && split) {
#define STRATA_TCN(NPV) gdn_tc_scan_kernel<true, true, NPV><<<HV, GTC_CW * 32, kTcScanSmemSplitNorm, s>>>(state, h, sc.ti, sc.te, sc.sc, sc.kq, sc.pi, y, t0, nch, T, nm->z, nm->gamma, nm->eps, nm->y16, nm->ld16)
if (np == 2) STRATA_TCN(2);
else if (np == 1) STRATA_TCN(1);
@@ -2915,12 +2936,13 @@ static void gdn_recurrence_impl(int variant, float* state, const float* h, const
// the split tensor-core scan can write y16 itself (STRATA_GDN_TC_FUSE=0 off): then the norm kernel below is skipped
static const bool tc_fuse = [] { const char* v = std::getenv("STRATA_GDN_TC_FUSE"); return v == nullptr || std::atoi(v) != 0; }();
const GdnTcNorm tcn = {z, gamma, eps, y16, ld16};
const GdnTcNorm* tcnp = with_out && tc_fuse ? &tcn : nullptr;
const int tcl = gdn_tc_tiles();
const GdnTcNorm* tcnp = with_out && tc_fuse && tcl == 1 ? &tcn : nullptr; // 2 tiles a head: the norm kernel (a head's norm spans both)
if (force != kGdnAuto) {
if (force == kGdnTc || force == kGdnTcSplit || force == kGdnTcS2 || force == kGdnTcS1) {
const bool sp = force != kGdnTc;
const int np = force == kGdnTcS2 ? 2 : force == kGdnTcS1 ? 1 : 3;
cudaError_t e = T >= kGdnTcMin ? gdn_rec_tc(state, h, gate, beta, y, T, cs, sp, np, sp ? tcnp : nullptr) : cudaErrorInvalidValue;
cudaError_t e = T >= kGdnTcMin ? gdn_rec_tc(state, h, gate, beta, y, T, cs, sp, np, sp ? tcnp : nullptr, sp ? tcl : 1) : cudaErrorInvalidValue;
if (e == cudaSuccess) { sel = force; normed = sp && tcnp != nullptr; }
else why = T < kGdnTcMin ? "T below the tc minimum (128)" : gdn_chunk_why(true, e, false, sp);
} else if (force == kGdnChunked) {
@@ -2948,7 +2970,7 @@ static void gdn_recurrence_impl(int variant, float* state, const float* h, const
if (pr) why = gdn_chunk_why(true, e, false);
// STRATA_GDN_TC_NP=1|2 (experiment, not a default): fewer split passes a product on this plan, other bits
static const int tc_np = [] { const char* v = std::getenv("STRATA_GDN_TC_NP"); const int x = v ? std::atoi(v) : 3; return x == 1 || x == 2 ? x : 3; }();
e = gdn_rec_tc(state, h, gate, beta, y, T, cs, true, tc_np, tcnp);
e = gdn_rec_tc(state, h, gate, beta, y, T, cs, true, tc_np, tcnp, tcl);
if (e == cudaSuccess) { sel = kGdnTcSplit; why = ""; normed = tcnp != nullptr; }
}
}