mirror of
https://github.com/Niko1221/Strata.git
synced 2026-11-07 06:15:58 -10:00
GDN: production-linked parity/bench harness, dispatch probe, FP64 fixtures, contract doc; opt-in STRATA_GDN_TC on 99-KB cards (split v/T/P buffering)
This commit is contained in:
Niko1221 committed
1 parent
870aedd9e1
commit
06f212aefe
7 files changed
+1030
-47
No files matched your search
@@ -1278,6 +1278,20 @@ if(STRATA_ENABLE_CUDA OR STRATA_ENABLE_HIP)
|
||||
# "not on this card" for the tc variant
|
||||
add_test(NAME gdn_rec_parity_selftest_frag COMMAND gdn_rec_parity --selftest-frag)
|
||||
set_tests_properties(gdn_rec_parity_selftest_frag PROPERTIES SKIP_RETURN_CODE 77)
|
||||
# the production-linked GDN parity / benchmark harness (no kernel copies: it calls gdn_gates, gdn_conv and
|
||||
# gdn_recurrence of strata_prefill through the GdnProbe of include/strata/prefill/gdn_probe.hpp, against an FP64
|
||||
# oracle from tests/gdn/gdn_fixture.hpp). Usage: docs/GDN_CONTRACT.md. The smoke test is small and CPU-cheap.
|
||||
execute_process(COMMAND git -C ${CMAKE_CURRENT_SOURCE_DIR} rev-parse --short=9 HEAD OUTPUT_VARIABLE _gdn_sha
|
||||
OUTPUT_STRIP_TRAILING_WHITESPACE ERROR_QUIET)
|
||||
if(NOT _gdn_sha)
|
||||
set(_gdn_sha unknown)
|
||||
endif()
|
||||
add_executable(gdn_harness src/prefill/gdn_harness.cu)
|
||||
target_include_directories(gdn_harness PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/tests/gdn)
|
||||
target_compile_definitions(gdn_harness PRIVATE GDN_HARNESS_SHA="${_gdn_sha}")
|
||||
target_link_libraries(gdn_harness PRIVATE strata_prefill)
|
||||
add_test(NAME gdn_harness_smoke COMMAND gdn_harness --T 400 --chunk 160 --cuts 37 --paths auto,keyhead,pipe --ref-heads 0,17)
|
||||
set_tests_properties(gdn_harness_smoke PROPERTIES TIMEOUT 300)
|
||||
# the prompt path's BF16 products against cuBLAS's own BF16 GEMM, within fp32 rounding (sm_7x runs them on the
|
||||
# FP16 tensor cores; GPU, synthetic, no model; --bench adds timings)
|
||||
add_executable(gemm_bf16_parity src/prefill/gemm_bf16_parity.cu)
|
||||
|
||||
@@ -0,0 +1,88 @@
|
||||
# GDN prompt-path contract (Flash-Next, CUDA reference)
|
||||
|
||||
Audited on main `870aedd9e` (v0.2.8), `src/prefill/kernels.cu` and `src/prefill/prefill.cpp`. Everything below is read from
|
||||
that source; the harness (`gdn_harness`, `tests/gdn/gdn_fixture.hpp`) checks it against an FP64 oracle that starts from the
|
||||
raw inputs, so a wrong reading shows up as a failed parity case.
|
||||
|
||||
## Shapes
|
||||
|
||||
`S = 128` (head size), `HK = 16` key heads, `HV = 48` value heads (3 per key head, `VPK = 3`), `C = 2*HK*S + HV*S = 10240`
|
||||
channels per token. Value head `vh` uses key head `qh = vh % HK` (NOT `vh / 3`; FLA groups the other way, an adapter is needed
|
||||
for any comparison with FLA). One layer's prompt call has `T` tokens (the engine's prefill chunk, not the whole prompt); the
|
||||
state and the convolution history carry from call to call.
|
||||
|
||||
## One layer, in the order the engine runs it (`prefill.cpp`, `pt.mark` stages)
|
||||
|
||||
| Stage (STRATA_PREFILL_TIMING name) | Work | In the harness |
|
||||
|---|---|---|
|
||||
| `gdn` (kPfGdn) | attn_qkv, attn_gate (z), ssm_alpha, ssm_beta projections (quantized GEMMs) | no (not a GDN kernel) |
|
||||
| `gdn conv+gates` (kPfGdnConv) | `gdn_gates`, `gdn_conv` | yes, timed separately |
|
||||
| `gdn recurrence` (kPfGdnRec) | `gdn_recurrence`: scan + output norm | yes |
|
||||
| `gdn out proj` (kPfGdnOut) | ssm_out projection from `y16` (quantized GEMM) | no |
|
||||
|
||||
Note the stage names: the profile line `gdn` is the INPUT projections, `gdn out proj` is the output projection. The task's
|
||||
"GDN stage = kPfGdn + Conv + Rec + Out" therefore contains two dense GEMM stages the recurrence work cannot speed up.
|
||||
|
||||
## Tensors (all FP32 unless noted, row-major, `t` = token in the call)
|
||||
|
||||
| Tensor | Layout | Meaning |
|
||||
|---|---|---|
|
||||
| `qkv` | `[T][C]` | projection output, channels `q (HK*S) | k (HK*S) | v (HV*S)` |
|
||||
| `ab` | `[T][2*HV]` | alpha at `[t][h]`, beta input at `[t][HV+h]` |
|
||||
| `dt`, `ssm_a` | `[HV]` | per-head bias and (negative) decay base |
|
||||
| `conv_w` | `[C][4]` | depthwise causal 4-tap filter, tap 3 is the current token |
|
||||
| `history` | `[C][3]` | the three previous raw `qkv` inputs of every channel (a state object of its own) |
|
||||
| `state` | `[S][HV][S]` (row i, head, column) | recurrent state, in place; `history` follows it in the engine's layer state |
|
||||
| `gate`, `beta` | `[T][HV]` | outputs of `gdn_gates` |
|
||||
| `h` | `[T][C]` | outputs of `gdn_conv` (SiLU, then q and k L2-normalized per head) |
|
||||
| `z` | `[T][HV*S]` | output gate input |
|
||||
| `y` | `[T][HV*S]` FP32 | scratch the recurrence writes the raw output to |
|
||||
| `y16` | FP16, row stride `ld16` | normalized output the out projection reads |
|
||||
|
||||
`ld16 = HV*S = 6144`, or `6144 + 64` when STRATA_PF_PAD is on and `T >= max(pf_switch_min_t, 64)`. A kernel must write exactly
|
||||
the 6144 columns and leave the pad columns untouched (the harness fills them with a sentinel and fails if they change).
|
||||
|
||||
## Math (per value head, `g_t = exp(gate_t)`)
|
||||
|
||||
- `gate_t = softplus(alpha_t + dt) * ssm_a` (softplus is the identity above 20); `beta_t = sigmoid(beta_input_t)`
|
||||
- `h_c(t) = silu(sum_k w[c][k] * x_c(t-3+k))` with `x_c(t<0)` from `history`; q and k heads are then scaled by `1/sqrt(sum^2 + 1e-6)`
|
||||
- `kv = S^T k`, `delta = (v - g*kv) * beta`, `S <- g*S + k delta^T`, `o = S^T q / sqrt(128)` (uses the UPDATED state)
|
||||
- `y16 = fp16( o * rsqrt(mean(o^2) + 1e-6) * gamma[col] * sigmoid(z) )`
|
||||
- after the call: `history <- last three raw inputs` (older history where `T < 3`)
|
||||
|
||||
## State mutation rules
|
||||
|
||||
- A call with `T <= 0` does nothing.
|
||||
- The call advances `state` and `history` in place. A path that cannot run must decline BEFORE touching the state (the helper
|
||||
returns an error and "nothing was launched"); once a kernel has started, there is no replay. A forced path that is not
|
||||
eligible launches nothing and reports kGdnNone with the reason (`gdn_probe.hpp`); the harness never times it.
|
||||
- Calls are sequential per layer state. The chunked path's scratch is one static buffer per device (single prompt stream per
|
||||
device); concurrent prompt streams on one device are not supported by it.
|
||||
|
||||
## Dispatch on CUDA (production, no probe)
|
||||
|
||||
1. `STRATA_GDN_TC=1` and `T >= 128`: `gdn_rec_tc` (tensor-core scan; sm_80+, >= 48 SMs, ~122 KB smem per block). Not mentioned in the
|
||||
2026-10-26 research pack: it landed with PR #3779. FP16 operands, FP32 accumulation: mixed-precision tier.
|
||||
2. `STRATA_GDN_CHUNKED=1` (`=2` skips the SM-count test) and `T >= 128`: `gdn_rec_chunked` (FP32, C = 32, scan grid `HK x S/16 = 128`
|
||||
blocks of 384 threads, 93.7 KB smem; sm_80+).
|
||||
3. else `gdn_rec_kh_kernel` when `gdn_keyhead_ok()` (sm_80+, all 64 blocks resident), otherwise `gdn_rec_cols_pipe_kernel`
|
||||
(`STRATA_GDN_PIPELINE=0`: `gdn_rec_cols_kernel`).
|
||||
4. `gdn_out_norm_kernel` after any of them (it is the part that writes `y16`).
|
||||
`STRATA_GDN_REC_WARP` and `STRATA_GDN_CHUNK` exist only in `STRATA_EXPERIMENTAL_SM60` builds and act only on a cc 7.0 device at run
|
||||
time (`gdn_chunk_available()`): they do nothing on sm_60, sm_86 or sm_120.
|
||||
|
||||
## Numerical tiers (gdn-research 06)
|
||||
|
||||
1. exact: output and state bitwise equal to the incumbent (`kh` and `pipe` are the same bits by construction).
|
||||
2. reordered FP32 (chunked): same recurrence, other summation order. Verdict in the harness: FP64-oracle error not above
|
||||
`max(floor, 3 x the incumbent's own error)` with floors 2e-3 (normalized max of the FP16 output) and 1e-4 (state); finite; the
|
||||
conv history exact.
|
||||
3. mixed precision (tc): same checks, plus the stress fixtures (`--stress`).
|
||||
The harness prints `normalized_max = max|d| / max|ref|` and relative L2 against the oracle and against the incumbent over ALL
|
||||
heads, tokens and state elements.
|
||||
|
||||
## Test hook
|
||||
|
||||
`include/strata/prefill/gdn_probe.hpp`: a thread-local `GdnProbe` installed by a test selects/forces a path, a smaller chunked scratch
|
||||
segment (`gsb`), or the SM-count bypass, and reads back `selected`, `wanted` and the fallback `reason`. With no probe installed the
|
||||
dispatch is unchanged. The engine never installs one.
|
||||
@@ -0,0 +1,56 @@
|
||||
// include/strata/prefill/gdn_probe.hpp - the GDN recurrence dispatch, made observable and forceable by tests.
|
||||
//
|
||||
// strata::prefill::gdn_recurrence picks its kernel from environment switches that are read once per process. A
|
||||
// test or benchmark that links the real strata_prefill library installs a GdnProbe on its thread to
|
||||
// * read back which path ran for the last call (`selected`) and why a path that was asked for did not run
|
||||
// (`reason`), instead of guessing from timings, and
|
||||
// * force one path for the call (`force`): a forced path that is not eligible launches NOTHING and reports
|
||||
// kGdnNone plus the reason, so a harness can never time a silent fallback as if it were the candidate.
|
||||
// Without a probe (the engine) the dispatch is byte-for-byte what it was: one thread-local pointer read.
|
||||
// Header only; CUDA builds fill the probe (src/prefill/kernels.cu), other backends may adopt the same struct.
|
||||
#pragma once
|
||||
|
||||
namespace strata::prefill {
|
||||
|
||||
enum GdnPath : int {
|
||||
kGdnNone = -1, // out: nothing was launched
|
||||
kGdnAuto = 0, // in: the production choice (the environment switches decide)
|
||||
kGdnKeyHead = 1, // gdn_rec_kh_kernel: the three value heads of a key head in one thread (sm_80+, 64 blocks resident)
|
||||
kGdnPipe = 2, // gdn_rec_cols_pipe_kernel: the software-pipelined column kernel
|
||||
kGdnCols = 3, // gdn_rec_cols_kernel (STRATA_GDN_PIPELINE=0)
|
||||
kGdnChunked = 4, // STRATA_GDN_CHUNKED: FP32 chunked prep + scan (CUDA, C = 32)
|
||||
kGdnTc = 5, // STRATA_GDN_TC=1: chunked, scan on the tensor cores (CUDA sm_80+, 125 KB of shared memory a block)
|
||||
kGdnTcSplit = 6, // the same scan with v / T / P single-buffered: 98 KB, fits the 99-KB cards (sm_86, sm_89, sm_120)
|
||||
kGdnPathCount = 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 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
|
||||
int wanted = kGdnAuto; // the path that was asked for: `force`, or the first the environment enabled
|
||||
const char* reason = ""; // when selected != wanted: why the wanted path did not run
|
||||
long calls = 0;
|
||||
};
|
||||
|
||||
inline thread_local GdnProbe* g_gdn_probe = nullptr;
|
||||
inline void gdn_probe_install(GdnProbe* p) { g_gdn_probe = p; }
|
||||
inline GdnProbe* gdn_probe_get() { return g_gdn_probe; }
|
||||
|
||||
inline const char* gdn_path_name(int p) {
|
||||
switch (p) {
|
||||
case kGdnKeyHead: return "keyhead";
|
||||
case kGdnPipe: return "pipe";
|
||||
case kGdnCols: return "cols";
|
||||
case kGdnChunked: return "chunked";
|
||||
case kGdnTc: return "tc";
|
||||
case kGdnTcSplit: return "tcs";
|
||||
case kGdnAuto: return "auto";
|
||||
default: return "none";
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace strata::prefill
|
||||
@@ -0,0 +1,470 @@
|
||||
// src/prefill/gdn_harness.cu - production-linked parity and benchmark harness for the prompt GDN path (CUDA).
|
||||
//
|
||||
// Links the real strata_prefill library and calls the real gdn_gates / gdn_conv / gdn_recurrence: there is not one
|
||||
// copy of a kernel in this file. Which recurrence runs is chosen through the GdnProbe (gdn_probe.hpp): a path can be
|
||||
// forced, the path that actually ran and the reason a wanted path did not run are read back after every call, and a
|
||||
// forced path that is not eligible launches nothing and is reported INELIGIBLE (it is never timed as a fallback).
|
||||
// Fixtures and the FP64 oracle are tests/gdn/gdn_fixture.hpp (shared with the other backends' harnesses).
|
||||
// Usage and output format: Desktop\30thOctober\results\gdn\HARNESS.md and docs/GDN_CONTRACT.md.
|
||||
#include "gdn_fixture.hpp"
|
||||
#include "strata/prefill/gdn_probe.hpp"
|
||||
#include "strata/prefill/kernels.hpp"
|
||||
|
||||
#include <cuda_runtime.h>
|
||||
|
||||
#include <algorithm>
|
||||
#include <chrono>
|
||||
#include <cstdio>
|
||||
#include <cstdlib>
|
||||
#include <cstring>
|
||||
#include <map>
|
||||
#include <set>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
#ifndef GDN_HARNESS_SHA
|
||||
#define GDN_HARNESS_SHA "unknown"
|
||||
#endif
|
||||
|
||||
using namespace gdnfx;
|
||||
namespace sp = strata::prefill;
|
||||
|
||||
namespace {
|
||||
|
||||
void ck(cudaError_t e, const char* what) {
|
||||
if (e != cudaSuccess) {
|
||||
std::fprintf(stderr, "%s: %s\n", what, cudaGetErrorString(e));
|
||||
std::exit(2);
|
||||
}
|
||||
}
|
||||
template <typename T> T* dalloc(size_t n) {
|
||||
T* p = nullptr;
|
||||
ck(cudaMalloc(&p, n * sizeof(T)), "cudaMalloc");
|
||||
return p;
|
||||
}
|
||||
template <typename T> T* dupload(const std::vector<T>& v) {
|
||||
T* p = dalloc<T>(v.size());
|
||||
ck(cudaMemcpy(p, v.data(), v.size() * sizeof(T), cudaMemcpyHostToDevice), "upload");
|
||||
return p;
|
||||
}
|
||||
|
||||
struct Args {
|
||||
int64_t T = 4096;
|
||||
std::vector<int64_t> Ts; // --T a,b,c : several lengths (one case each)
|
||||
int64_t chunk = 2048; // tokens per engine call
|
||||
std::vector<int64_t> cuts; // explicit first call sizes (unaligned splits), then `chunk`
|
||||
std::vector<int> paths; // default: auto + every forced path
|
||||
int incumbent = sp::kGdnKeyHead; // the path others are compared with (falls to pipe where ineligible)
|
||||
bool parity = true, bench = false, ref = true;
|
||||
std::vector<int> ref_heads = {0, 17, 34, 5, 22, 47};
|
||||
int reps = 7, warmup = 2;
|
||||
bool any_sms = false;
|
||||
int gsb = 0;
|
||||
bool fallback_incumbent = false; // a forced path takes the incumbent for calls it cannot run (mirrors production)
|
||||
int64_t ld16 = 0;
|
||||
FixtureSpec fx;
|
||||
std::string json;
|
||||
std::string label;
|
||||
bool stress = false; // run the stress fixtures (parity only)
|
||||
};
|
||||
|
||||
std::vector<int64_t> parse_list(const char* s) {
|
||||
std::vector<int64_t> v;
|
||||
for (const char* p = s; *p;) {
|
||||
char* e = nullptr;
|
||||
long long x = std::strtoll(p, &e, 10);
|
||||
if (e == p) break;
|
||||
if (*e == 'k' || *e == 'K') { x *= 1024; ++e; }
|
||||
v.push_back(x);
|
||||
p = *e == ',' ? e + 1 : e;
|
||||
}
|
||||
return v;
|
||||
}
|
||||
int path_from_name(const std::string& n) {
|
||||
for (int p = 0; p < sp::kGdnPathCount; ++p)
|
||||
if (n == sp::gdn_path_name(p)) return p;
|
||||
return -2;
|
||||
}
|
||||
|
||||
struct Dev { // device-resident fixture
|
||||
int64_t T = 0;
|
||||
float *qkv = nullptr, *ab = nullptr, *dt = nullptr, *ssm_a = nullptr, *conv_w = nullptr, *z = nullptr, *gamma = nullptr;
|
||||
float* state0 = nullptr; // [S*HV*S state | C*3 history]
|
||||
float *gate = nullptr, *beta = nullptr, *h = nullptr, *y = nullptr, *state = nullptr;
|
||||
uint16_t* y16 = nullptr;
|
||||
int64_t ld16 = 0, chunk_max = 0;
|
||||
char* flush = nullptr;
|
||||
size_t flush_bytes = 256u << 20;
|
||||
};
|
||||
constexpr size_t kStateFloats = (size_t) S * HV * S;
|
||||
|
||||
Dev upload(const Fixture& f, int64_t chunk_max, int64_t ld16) {
|
||||
Dev d;
|
||||
d.T = f.T;
|
||||
d.chunk_max = chunk_max;
|
||||
d.ld16 = ld16 > 0 ? ld16 : (int64_t) HV * S;
|
||||
d.qkv = dupload(f.qkv);
|
||||
d.ab = dupload(f.ab);
|
||||
d.dt = dupload(f.dt);
|
||||
d.ssm_a = dupload(f.ssm_a);
|
||||
d.conv_w = dupload(f.conv_w);
|
||||
d.z = dupload(f.z);
|
||||
d.gamma = dupload(f.gamma);
|
||||
std::vector<float> st(kStateFloats + (size_t) C * 3);
|
||||
std::memcpy(st.data(), f.state0.data(), kStateFloats * 4);
|
||||
std::memcpy(st.data() + kStateFloats, f.hist.data(), (size_t) C * 3 * 4);
|
||||
d.state0 = dupload(st);
|
||||
d.state = dalloc<float>(st.size());
|
||||
d.gate = dalloc<float>((size_t) chunk_max * HV);
|
||||
d.beta = dalloc<float>((size_t) chunk_max * HV);
|
||||
d.h = dalloc<float>((size_t) chunk_max * C);
|
||||
d.y = dalloc<float>((size_t) chunk_max * HV * S);
|
||||
d.y16 = dalloc<uint16_t>((size_t) f.T * d.ld16 + 64);
|
||||
ck(cudaMemset(d.y16, 0xAB, ((size_t) f.T * d.ld16 + 64) * 2), "memset y16 sentinel");
|
||||
d.flush = dalloc<char>(d.flush_bytes);
|
||||
return d;
|
||||
}
|
||||
void release(Dev& d) {
|
||||
for (void* p : {(void*) d.qkv, (void*) d.ab, (void*) d.dt, (void*) d.ssm_a, (void*) d.conv_w, (void*) d.z, (void*) d.gamma,
|
||||
(void*) d.state0, (void*) d.state, (void*) d.gate, (void*) d.beta, (void*) d.h, (void*) d.y,
|
||||
(void*) d.y16, (void*) d.flush})
|
||||
cudaFree(p);
|
||||
d = Dev();
|
||||
}
|
||||
|
||||
std::vector<int64_t> plan_calls(int64_t T, const Args& a) {
|
||||
std::vector<int64_t> v;
|
||||
int64_t left = T;
|
||||
for (int64_t c : a.cuts) {
|
||||
if (left <= 0) break;
|
||||
v.push_back(std::min(c, left));
|
||||
left -= v.back();
|
||||
}
|
||||
while (left > 0) { v.push_back(std::min(a.chunk, left)); left -= v.back(); }
|
||||
return v;
|
||||
}
|
||||
|
||||
struct RunOut {
|
||||
int path = 0;
|
||||
bool eligible = true;
|
||||
std::string reason;
|
||||
std::set<std::string> logged; // distinct "wanted->selected [reason]" over the calls
|
||||
int fallback_calls = 0;
|
||||
double ms_gates = 0, ms_conv = 0, ms_rec = 0;
|
||||
std::vector<uint16_t> y16; // dense [T][HV*S]
|
||||
std::vector<float> state;
|
||||
};
|
||||
|
||||
// Runs the layer pipeline for every call of the plan. timed: events around each stage.
|
||||
RunOut run_path(Dev& d, const Args& a, const std::vector<int64_t>& calls, int path, bool timed, bool download) {
|
||||
RunOut r;
|
||||
r.path = path;
|
||||
ck(cudaMemcpy(d.state, d.state0, (kStateFloats + (size_t) C * 3) * 4, cudaMemcpyDeviceToDevice), "reset state");
|
||||
float* conv_hist = d.state + kStateFloats;
|
||||
sp::GdnProbe probe;
|
||||
probe.any_sms = a.any_sms;
|
||||
probe.gsb = a.gsb;
|
||||
sp::gdn_probe_install(&probe);
|
||||
std::vector<cudaEvent_t> ev;
|
||||
if (timed) {
|
||||
ev.resize(calls.size() * 4);
|
||||
for (auto& e : ev) ck(cudaEventCreate(&e), "event");
|
||||
}
|
||||
int64_t t0 = 0;
|
||||
size_t ci = 0;
|
||||
for (int64_t T : calls) {
|
||||
if (timed) ck(cudaEventRecord(ev[ci * 4]), "ev");
|
||||
sp::gdn_gates(d.ab + t0 * 2 * HV, d.dt, d.ssm_a, d.gate, d.beta, T, nullptr);
|
||||
if (timed) ck(cudaEventRecord(ev[ci * 4 + 1]), "ev");
|
||||
sp::gdn_conv(conv_hist, d.qkv + t0 * C, d.conv_w, d.h, T, EPS, nullptr);
|
||||
if (timed) ck(cudaEventRecord(ev[ci * 4 + 2]), "ev");
|
||||
probe.force = path;
|
||||
sp::gdn_recurrence(d.state, d.h, d.gate, d.beta, d.z + t0 * HV * S, d.gamma, EPS, d.y, d.y16 + t0 * d.ld16, T, nullptr, d.ld16);
|
||||
if (probe.selected == sp::kGdnNone && a.fallback_incumbent && path != a.incumbent) {
|
||||
probe.force = a.incumbent;
|
||||
sp::gdn_recurrence(d.state, d.h, d.gate, d.beta, d.z + t0 * HV * S, d.gamma, EPS, d.y, d.y16 + t0 * d.ld16, T, nullptr, d.ld16);
|
||||
if (probe.selected == sp::kGdnNone) probe.force = sp::kGdnPipe, sp::gdn_recurrence(d.state, d.h, d.gate, d.beta, d.z + t0 * HV * S, d.gamma, EPS, d.y, d.y16 + t0 * d.ld16, T, nullptr, d.ld16);
|
||||
++r.fallback_calls;
|
||||
}
|
||||
if (timed) ck(cudaEventRecord(ev[ci * 4 + 3]), "ev");
|
||||
if (probe.selected == sp::kGdnNone) {
|
||||
r.eligible = false;
|
||||
r.reason = std::string("call of ") + std::to_string((long long) T) + " tokens: " + probe.reason;
|
||||
break;
|
||||
}
|
||||
char b[200];
|
||||
std::snprintf(b, sizeof b, "%s->%s%s%s", sp::gdn_path_name(probe.wanted), sp::gdn_path_name(probe.selected),
|
||||
probe.reason[0] ? " [" : "", probe.reason[0] ? probe.reason : "");
|
||||
std::string s = b;
|
||||
if (probe.reason[0]) s += "]";
|
||||
r.logged.insert(s);
|
||||
t0 += T;
|
||||
++ci;
|
||||
}
|
||||
ck(cudaDeviceSynchronize(), "run_path sync");
|
||||
sp::gdn_probe_install(nullptr);
|
||||
if (timed && r.eligible)
|
||||
for (size_t i = 0; i < ci; ++i) {
|
||||
float x = 0;
|
||||
cudaEventElapsedTime(&x, ev[i * 4], ev[i * 4 + 1]); r.ms_gates += x;
|
||||
cudaEventElapsedTime(&x, ev[i * 4 + 1], ev[i * 4 + 2]); r.ms_conv += x;
|
||||
cudaEventElapsedTime(&x, ev[i * 4 + 2], ev[i * 4 + 3]); r.ms_rec += x;
|
||||
}
|
||||
for (auto e : ev) cudaEventDestroy(e);
|
||||
if (download && r.eligible) {
|
||||
r.y16.resize((size_t) d.T * HV * S);
|
||||
if (d.ld16 == (int64_t) HV * S) {
|
||||
ck(cudaMemcpy(r.y16.data(), d.y16, r.y16.size() * 2, cudaMemcpyDeviceToHost), "y16 down");
|
||||
} else {
|
||||
ck(cudaMemcpy2D(r.y16.data(), (size_t) HV * S * 2, d.y16, (size_t) d.ld16 * 2, (size_t) HV * S * 2, d.T, cudaMemcpyDeviceToHost), "y16 down 2d");
|
||||
// padding columns of a padded stride must be untouched (sentinel 0xABAB)
|
||||
std::vector<uint16_t> pad((size_t) d.T * d.ld16);
|
||||
ck(cudaMemcpy(pad.data(), d.y16, pad.size() * 2, cudaMemcpyDeviceToHost), "pad down");
|
||||
for (int64_t t = 0; t < d.T; ++t)
|
||||
for (int64_t c = HV * S; c < d.ld16; ++c)
|
||||
if (pad[(size_t) t * d.ld16 + c] != 0xABAB) { r.eligible = false; r.reason = "PADDING OVERWRITTEN (stride contract broken)"; return r; }
|
||||
}
|
||||
r.state.resize(kStateFloats + (size_t) C * 3);
|
||||
ck(cudaMemcpy(r.state.data(), d.state, r.state.size() * 4, cudaMemcpyDeviceToHost), "state down");
|
||||
}
|
||||
return r;
|
||||
}
|
||||
|
||||
double median(std::vector<double> v) {
|
||||
std::sort(v.begin(), v.end());
|
||||
return v.empty() ? 0 : v[v.size() / 2];
|
||||
}
|
||||
|
||||
struct Cmp { double y_nmax = 0, y_rl2 = 0, s_nmax = 0, s_rl2 = 0; long nonfinite = 0; };
|
||||
|
||||
Cmp vs_ref(const RunOut& r, const RefResult& ref, int64_t T) {
|
||||
Cmp c;
|
||||
ErrStat ey, es;
|
||||
for (size_t n = 0; n < ref.heads.size(); ++n) {
|
||||
const int vh = ref.heads[n];
|
||||
for (int64_t t = 0; t < T; ++t)
|
||||
for (int col = 0; col < S; ++col)
|
||||
ey.add(half_to_float(r.y16[(size_t) t * HV * S + (size_t) vh * S + col]), ref.y[n][(size_t) t * S + col]);
|
||||
for (int i = 0; i < S; ++i)
|
||||
for (int col = 0; col < S; ++col) es.add(r.state[((size_t) i * HV + vh) * S + col], ref.state[n][(size_t) i * S + col]);
|
||||
}
|
||||
c.y_nmax = ey.normalized_max(); c.y_rl2 = ey.rel_l2(); c.s_nmax = es.normalized_max(); c.s_rl2 = es.rel_l2();
|
||||
c.nonfinite = ey.nonfinite + es.nonfinite;
|
||||
return c;
|
||||
}
|
||||
// candidate against the incumbent over every head, token and state element (not a subset)
|
||||
Cmp vs_incumbent(const RunOut& r, const RunOut& inc, int64_t T) {
|
||||
Cmp c;
|
||||
ErrStat ey, es;
|
||||
for (size_t i = 0; i < r.y16.size(); ++i) ey.add(half_to_float(r.y16[i]), half_to_float(inc.y16[i]));
|
||||
for (size_t i = 0; i < kStateFloats; ++i) es.add(r.state[i], inc.state[i]);
|
||||
c.y_nmax = ey.normalized_max(); c.y_rl2 = ey.rel_l2(); c.s_nmax = es.normalized_max(); c.s_rl2 = es.rel_l2();
|
||||
c.nonfinite = ey.nonfinite + es.nonfinite;
|
||||
(void) T;
|
||||
return c;
|
||||
}
|
||||
|
||||
std::string join(const std::set<std::string>& s) {
|
||||
std::string o;
|
||||
for (auto& x : s) { if (!o.empty()) o += "; "; o += x; }
|
||||
return o;
|
||||
}
|
||||
|
||||
void json_line(const Args& a, const char* kind, const char* fixture, int64_t T, const RunOut& r, const Cmp* ref, const Cmp* inc,
|
||||
bool exact, double ms_med, double ms_min, const char* verdict, const cudaDeviceProp& pr) {
|
||||
if (a.json.empty()) return;
|
||||
FILE* f = std::fopen(a.json.c_str(), "a");
|
||||
if (!f) return;
|
||||
std::fprintf(f, "{\"kind\":\"%s\",\"label\":\"%s\",\"device\":\"%s\",\"sms\":%d,\"cc\":\"%d.%d\",\"sha\":\"%s\",\"fixture\":\"%s\","
|
||||
"\"T\":%lld,\"chunk\":%lld,\"path\":\"%s\",\"eligible\":%s,\"reason\":\"%s\",\"log\":\"%s\",\"fallback_calls\":%d,"
|
||||
"\"ms_gates\":%.4f,\"ms_conv\":%.4f,\"ms_rec\":%.4f,\"ms_total_med\":%.4f,\"ms_total_min\":%.4f,\"exact_vs_incumbent\":%s,\"verdict\":\"%s\"",
|
||||
kind, a.label.c_str(), pr.name, pr.multiProcessorCount, pr.major, pr.minor, GDN_HARNESS_SHA, fixture, (long long) T,
|
||||
(long long) a.chunk, sp::gdn_path_name(r.path), r.eligible ? "true" : "false", r.reason.c_str(), join(r.logged).c_str(),
|
||||
r.fallback_calls, r.ms_gates, r.ms_conv, r.ms_rec, ms_med, ms_min, exact ? "true" : "false", verdict);
|
||||
if (ref) std::fprintf(f, ",\"ref_y_nmax\":%.3e,\"ref_y_rl2\":%.3e,\"ref_s_nmax\":%.3e,\"ref_s_rl2\":%.3e,\"nonfinite\":%ld", ref->y_nmax, ref->y_rl2, ref->s_nmax, ref->s_rl2, ref->nonfinite);
|
||||
if (inc) std::fprintf(f, ",\"inc_y_nmax\":%.3e,\"inc_y_rl2\":%.3e,\"inc_s_nmax\":%.3e,\"inc_s_rl2\":%.3e", inc->y_nmax, inc->y_rl2, inc->s_nmax, inc->s_rl2);
|
||||
std::fprintf(f, "}\n");
|
||||
std::fclose(f);
|
||||
}
|
||||
|
||||
// One case: a fixture of T tokens, every wanted path. Returns the number of failed verdicts.
|
||||
int run_case(const Args& a, const FixtureSpec& spec, const char* fname, const cudaDeviceProp& pr) {
|
||||
const int64_t T = spec.T;
|
||||
const std::vector<int64_t> calls = plan_calls(T, a);
|
||||
int64_t cmax = 0;
|
||||
for (auto c : calls) cmax = std::max(cmax, c);
|
||||
Fixture f = make_fixture(spec);
|
||||
Dev d = upload(f, cmax, a.ld16);
|
||||
std::printf("\n=== case %s T=%lld calls=%zu (first %lld, max %lld) ld16=%lld device %s (%d SMs, cc %d.%d)\n", fname, (long long) T,
|
||||
calls.size(), (long long) calls[0], (long long) cmax, (long long) d.ld16, pr.name, pr.multiProcessorCount, pr.major, pr.minor);
|
||||
RefResult ref;
|
||||
if (a.parity && a.ref) {
|
||||
auto t0 = std::chrono::steady_clock::now();
|
||||
ref = ref_run(f, a.ref_heads);
|
||||
std::printf("FP64 oracle for %zu value heads: %.1f s\n", a.ref_heads.size(),
|
||||
std::chrono::duration<double>(std::chrono::steady_clock::now() - t0).count());
|
||||
}
|
||||
// the incumbent first (it may itself be ineligible: then pipe, which always runs)
|
||||
std::vector<int> order = a.paths;
|
||||
RunOut inc = run_path(d, a, calls, a.incumbent, false, true);
|
||||
int inc_path = a.incumbent;
|
||||
if (!inc.eligible) {
|
||||
std::printf("incumbent %s INELIGIBLE (%s): using pipe\n", sp::gdn_path_name(a.incumbent), inc.reason.c_str());
|
||||
inc = run_path(d, a, calls, sp::kGdnPipe, false, true);
|
||||
inc_path = sp::kGdnPipe;
|
||||
}
|
||||
std::map<int, RunOut> runs;
|
||||
for (int p : order) {
|
||||
RunOut r = (p == inc_path) ? inc : run_path(d, a, calls, p, false, true);
|
||||
runs[p] = std::move(r);
|
||||
}
|
||||
int fails = 0;
|
||||
std::map<int, std::vector<double>> total_ms;
|
||||
std::map<int, RunOut> timed_last;
|
||||
if (a.bench) {
|
||||
for (int w = 0; w < a.warmup + a.reps; ++w)
|
||||
for (int p : order) {
|
||||
if (!runs[p].eligible) continue;
|
||||
ck(cudaMemset(d.flush, w & 255, d.flush_bytes), "flush");
|
||||
RunOut t = run_path(d, a, calls, p, true, false);
|
||||
if (w >= a.warmup) {
|
||||
total_ms[p].push_back(t.ms_gates + t.ms_conv + t.ms_rec);
|
||||
timed_last[p].ms_gates += t.ms_gates / a.reps;
|
||||
timed_last[p].ms_conv += t.ms_conv / a.reps;
|
||||
timed_last[p].ms_rec += t.ms_rec / a.reps;
|
||||
}
|
||||
}
|
||||
}
|
||||
std::printf("%-8s %-6s %-9s %9s %9s %9s %9s %-34s %s\n", "path", "elig", "bitwise", "ref_y", "ref_state", "inc_y", "inc_state", "ms gates/conv/rec (mean)", "ran");
|
||||
for (int p : order) {
|
||||
RunOut& r = runs[p];
|
||||
const char* name = sp::gdn_path_name(p);
|
||||
if (!r.eligible) {
|
||||
std::printf("%-8s %-6s %s\n", name, "NO", r.reason.c_str());
|
||||
json_line(a, "case", fname, T, r, nullptr, nullptr, false, 0, 0, "INELIGIBLE", pr);
|
||||
continue;
|
||||
}
|
||||
Cmp cr{}, ci{};
|
||||
const bool have_ref = a.parity && a.ref;
|
||||
if (have_ref) cr = vs_ref(r, ref, T);
|
||||
ci = vs_incumbent(r, inc, T);
|
||||
const bool exact = r.y16 == inc.y16 && std::memcmp(r.state.data(), inc.state.data(), (kStateFloats + C * 3) * 4) == 0;
|
||||
// conv history after the last call is a copy of inputs: it must match the oracle's exactly
|
||||
bool hist_ok = true;
|
||||
if (have_ref) hist_ok = std::memcmp(r.state.data() + kStateFloats, ref.hist.data(), (size_t) C * 3 * 4) == 0;
|
||||
// verdict: finite, history exact, and not worse than max(floor, 3 x the incumbent's own error vs the oracle)
|
||||
const char* verdict = "PASS";
|
||||
if (have_ref) {
|
||||
const Cmp ir = vs_ref(inc, ref, T);
|
||||
const double ys = std::max(2e-3, 3 * ir.y_nmax), ss = std::max(1e-4, 3 * ir.s_nmax);
|
||||
if (cr.nonfinite || !hist_ok) verdict = "FAIL-nonfinite-or-history";
|
||||
else if (cr.y_nmax > ys) verdict = "FAIL-output-vs-oracle";
|
||||
else if (cr.s_nmax > ss) verdict = "FAIL-state-vs-oracle";
|
||||
else if (exact) verdict = "PASS-bitwise";
|
||||
} else if (exact) verdict = "PASS-bitwise";
|
||||
else verdict = "NO-ORACLE";
|
||||
if (!std::strncmp(verdict, "FAIL", 4)) ++fails;
|
||||
double med = 0, mn = 0;
|
||||
if (a.bench && !total_ms[p].empty()) {
|
||||
med = median(total_ms[p]);
|
||||
mn = *std::min_element(total_ms[p].begin(), total_ms[p].end());
|
||||
r.ms_gates = timed_last[p].ms_gates; r.ms_conv = timed_last[p].ms_conv; r.ms_rec = timed_last[p].ms_rec;
|
||||
}
|
||||
char tb[120] = "";
|
||||
if (a.bench) std::snprintf(tb, sizeof tb, "%.3f/%.3f/%.3f (tot med %.3f min %.3f)", r.ms_gates, r.ms_conv, r.ms_rec, med, mn);
|
||||
std::printf("%-8s %-6s %-9s %9.2e %9.2e %9.2e %9.2e %-34s %s %s\n", name, "yes", exact ? "identical" : "differs", cr.y_nmax, cr.s_nmax,
|
||||
ci.y_nmax, ci.s_nmax, tb, join(r.logged).c_str(), verdict);
|
||||
if (r.fallback_calls) std::printf(" (%d call(s) ran on the incumbent: the forced path could not take them)\n", r.fallback_calls);
|
||||
json_line(a, "case", fname, T, r, have_ref ? &cr : nullptr, &ci, exact, med, mn, verdict, pr);
|
||||
}
|
||||
release(d);
|
||||
return fails;
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
int main(int argc, char** argv) {
|
||||
Args a;
|
||||
bool paths_given = false;
|
||||
for (int i = 1; i < argc; ++i) {
|
||||
std::string k = argv[i];
|
||||
auto val = [&]() -> const char* { if (i + 1 >= argc) { std::fprintf(stderr, "%s needs a value\n", k.c_str()); std::exit(2); } return argv[++i]; };
|
||||
if (k == "--T") { a.Ts = parse_list(val()); }
|
||||
else if (k == "--chunk") a.chunk = parse_list(val())[0];
|
||||
else if (k == "--cuts") a.cuts = parse_list(val());
|
||||
else if (k == "--paths") {
|
||||
paths_given = true;
|
||||
std::string s = val();
|
||||
for (size_t p = 0; p <= s.size();) {
|
||||
size_t e = s.find(',', p);
|
||||
if (e == std::string::npos) e = s.size();
|
||||
const int id = path_from_name(s.substr(p, e - p));
|
||||
if (id < 0) { std::fprintf(stderr, "unknown path %s (auto,keyhead,pipe,cols,chunked,tc,tcs)\n", s.substr(p, e - p).c_str()); return 2; }
|
||||
a.paths.push_back(id);
|
||||
p = e + 1;
|
||||
}
|
||||
}
|
||||
else if (k == "--incumbent") a.incumbent = path_from_name(val());
|
||||
else if (k == "--parity") a.parity = true;
|
||||
else if (k == "--no-parity") a.parity = false;
|
||||
else if (k == "--bench") a.bench = true;
|
||||
else if (k == "--no-ref") a.ref = false;
|
||||
else if (k == "--ref-heads") { a.ref_heads.clear(); for (auto x : parse_list(val())) a.ref_heads.push_back((int) x); }
|
||||
else if (k == "--reps") a.reps = std::atoi(val());
|
||||
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 == "--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];
|
||||
else if (k == "--v-scale") a.fx.v_scale = (float) std::atof(val());
|
||||
else if (k == "--gate-scale") a.fx.gate_scale = (float) std::atof(val());
|
||||
else if (k == "--coherent-keys") a.fx.coherent_keys = true;
|
||||
else if (k == "--alt-beta") a.fx.alt_beta = true;
|
||||
else if (k == "--tiny-k") a.fx.tiny_k = true;
|
||||
else if (k == "--stress") a.stress = true;
|
||||
else if (k == "--seed") a.fx.seed = (uint64_t) std::atoll(val());
|
||||
else if (k == "--json") a.json = val();
|
||||
else if (k == "--label") a.label = val();
|
||||
else {
|
||||
std::fprintf(stderr, "unknown option %s (see Desktop\\30thOctober\\results\\gdn\\HARNESS.md)\n", k.c_str());
|
||||
return 2;
|
||||
}
|
||||
}
|
||||
if (a.Ts.empty()) a.Ts = {a.T};
|
||||
if (!paths_given) a.paths = {sp::kGdnAuto, sp::kGdnKeyHead, sp::kGdnPipe, sp::kGdnChunked, sp::kGdnTc, sp::kGdnTcSplit};
|
||||
if (a.incumbent < 0) { std::fprintf(stderr, "bad --incumbent\n"); return 2; }
|
||||
int dev = 0;
|
||||
ck(cudaGetDevice(&dev), "device");
|
||||
cudaDeviceProp pr{};
|
||||
ck(cudaGetDeviceProperties(&pr, dev), "props");
|
||||
int drv = 0, rt = 0;
|
||||
cudaDriverGetVersion(&drv);
|
||||
cudaRuntimeGetVersion(&rt);
|
||||
std::printf("gdn_harness sha=%s device=%s sms=%d cc=%d.%d driver=%d runtime=%d smem_optin=%zu\n", GDN_HARNESS_SHA, pr.name,
|
||||
pr.multiProcessorCount, pr.major, pr.minor, drv, rt, (size_t) pr.sharedMemPerBlockOptin);
|
||||
for (const char* e : {"STRATA_GDN_CHUNKED", "STRATA_GDN_TC", "STRATA_GDN_KEYHEAD", "STRATA_GDN_PIPELINE", "STRATA_GDN_REC_WARP",
|
||||
"STRATA_GDN_CONVL2", "STRATA_GDN_CONV_SERIAL", "STRATA_GDN_REC_HEADS", "STRATA_PF_PAD"})
|
||||
if (const char* v = std::getenv(e)) std::printf("env %s=%s\n", e, v);
|
||||
int fails = 0;
|
||||
for (int64_t T : a.Ts) {
|
||||
FixtureSpec s = a.fx;
|
||||
s.T = T;
|
||||
fails += run_case(a, s, "default", pr);
|
||||
if (a.stress) {
|
||||
struct { const char* n; const char* st; float v, g; bool coh, alt, tiny; } cases[] = {
|
||||
{"state-zero", "zero", 1, 1, false, false, false}, {"state-big", "big", 1, 1, false, false, false},
|
||||
{"state-small", "small", 1, 1, false, false, false}, {"v-x50", "random", 50, 1, false, false, false},
|
||||
{"slow-gates", "random", 1, 0.005f, false, false, false}, {"fast-gates", "random", 1, 3, false, false, false},
|
||||
{"coherent-keys", "random", 1, 1, true, false, false}, {"beta-endpoints", "random", 1, 1, false, true, false},
|
||||
{"tiny-keys", "random", 1, 1, false, false, true}};
|
||||
for (auto& c : cases) {
|
||||
FixtureSpec t = a.fx;
|
||||
t.T = T; t.state = c.st; t.v_scale = c.v; t.gate_scale = c.g; t.coherent_keys = c.coh; t.alt_beta = c.alt; t.tiny_k = c.tiny;
|
||||
fails += run_case(a, t, c.n, pr);
|
||||
}
|
||||
}
|
||||
}
|
||||
std::printf("\n%s: %d failed verdict(s)\n", fails ? "FAIL" : "OK", fails);
|
||||
return fails ? 1 : 0;
|
||||
}
|
||||
+45
-30
@@ -46,6 +46,11 @@ template <bool NORM>
|
||||
constexpr size_t kTcScanSmem = 2 * (4 * GCH * GTC_KP * 2 + GCH * GTC_VP * 4 + (NORM ? GCH * GTC_VP * 4 : 0) +
|
||||
4 * GCH * GTC_PP * 2 + GCH * 16 + (NORM ? GCH * GTC_CW * 4 : 0)) + 16;
|
||||
static_assert(kTcScanSmem<false> == 124944 && kTcScanSmem<true> == 160784, "gdn_tc.cuh: the shared-memory plan");
|
||||
// SPLIT (the 99-KB cards: sm_86 / sm_89 / sm_120): only k / q hi / lo and the token scalars stay double-buffered; v, T and P
|
||||
// (one buffer) are staged for the next chunk after this chunk's last read of them, so that load overlaps the state update:
|
||||
// 97,808 B against the 101,376 B opt-in of those cards (the full plan needs 124,944)
|
||||
constexpr size_t kTcScanSmemSplit = 2 * (4 * GCH * GTC_KP * 2 + GCH * 16) + GCH * GTC_VP * 4 + 4 * GCH * GTC_PP * 2 + 16;
|
||||
static_assert(kTcScanSmemSplit == 97808, "gdn_tc.cuh: the split shared-memory plan");
|
||||
// the prep's scratch for a super-block (gdn_tc_scratch, one a device, 60.3 MB): the T images x 2^et as
|
||||
// [chunk][vh][2: hi, lo][GCH][GCH] f16 and their scales 2^-et as [chunk][vh], the token scalars as float4
|
||||
// [chunk][vh][GCH], the k and q images x 2^8 as [chunk][qh][4: k hi, k lo, q hi, q lo][GCH][S] and the P images x 2^8
|
||||
@@ -414,7 +419,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>
|
||||
template <bool NORM, bool SPLIT = false>
|
||||
__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,
|
||||
@@ -426,12 +431,13 @@ gdn_tc_scan_kernel(float* __restrict__ state, const float* __restrict__ h, const
|
||||
uint16_t* kl = kh + 2 * GCH * GTC_KP; // [2][GCH][GTC_KP] what hi left out (lo)
|
||||
uint16_t* qh = kl + 2 * GCH * GTC_KP; // q likewise
|
||||
uint16_t* ql = qh + 2 * GCH * GTC_KP;
|
||||
float* sv = reinterpret_cast<float*>(ql + 2 * GCH * GTC_KP); // [2][GCH][GTC_VP] v
|
||||
uint16_t* th = reinterpret_cast<uint16_t*>(sv + 2 * GCH * GTC_VP); // [2][GCH][GTC_PP] T x 2^et hi / lo
|
||||
uint16_t* tl = th + 2 * GCH * GTC_PP;
|
||||
uint16_t* ph = tl + 2 * GCH * GTC_PP; // [2][GCH][GTC_PP] P x 2^8 hi / lo
|
||||
uint16_t* pl = ph + 2 * GCH * GTC_PP;
|
||||
float4* sc = reinterpret_cast<float4*>(pl + 2 * GCH * GTC_PP); // [2][GCH] (beta, exp(gamma), exp(gL - gamma), gamma)
|
||||
constexpr int NV = SPLIT ? 1 : 2; // buffers of v, T and P
|
||||
float* sv = reinterpret_cast<float*>(ql + 2 * GCH * GTC_KP); // [NV][GCH][GTC_VP] v
|
||||
uint16_t* th = reinterpret_cast<uint16_t*>(sv + NV * GCH * GTC_VP); // [NV][GCH][GTC_PP] T x 2^et hi / lo
|
||||
uint16_t* tl = th + NV * GCH * GTC_PP;
|
||||
uint16_t* ph = tl + NV * GCH * GTC_PP; // [NV][GCH][GTC_PP] P x 2^8 hi / lo
|
||||
uint16_t* pl = ph + NV * GCH * GTC_PP;
|
||||
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
|
||||
const int vh = blockIdx.x, qhd = vh % HK;
|
||||
const int tid = threadIdx.x, warp = tid >> 5, lane = tid & 31, gid = lane >> 2, tig = lane & 3;
|
||||
@@ -439,32 +445,36 @@ gdn_tc_scan_kernel(float* __restrict__ state, const float* __restrict__ h, const
|
||||
// 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
|
||||
// 16 pieces: 8 a thread), v (4 a thread), the two T images (1), the two P images (1), the scalars (warp 0, 1),
|
||||
// T's scale (thread 0, 4 B)
|
||||
auto stage = [&](int ch, int b) {
|
||||
// parts: bit 0 = the k / q images, the scalars and T's scale (buffer b); bit 1 = v, T and P (buffer bv: b, or 0 in SPLIT)
|
||||
auto stage = [&](int ch, int b, int parts) {
|
||||
const int64_t c_t0 = t0 + (int64_t) ch * GCH;
|
||||
const int n = (int) min((int64_t) GCH, T - c_t0);
|
||||
const uint16_t* img = kqg + ((size_t) ch * HK + qhd) * 4 * GCH * S;
|
||||
const int bv = SPLIT ? 0 : b;
|
||||
if (parts & 1) {
|
||||
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;
|
||||
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));
|
||||
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;
|
||||
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));
|
||||
}
|
||||
if (tid < GCH) gdn_cp16(reinterpret_cast<float*>(sc + b * GCH + tid), reinterpret_cast<const float*>(scg + ((int64_t) ch * HV + vh) * GCH + tid));
|
||||
if (tid == 0) gdn_cp4(ste + b, teg + (int64_t) ch * HV + vh);
|
||||
}
|
||||
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;
|
||||
float* dst = sv + (b * GCH + t) * GTC_VP + r;
|
||||
if (t < n) gdn_cp16(dst, h + (c_t0 + t) * C + 2 * HK * S + vh * S + r);
|
||||
else *reinterpret_cast<float4*>(dst) = make_float4(0.0f, 0.0f, 0.0f, 0.0f);
|
||||
}
|
||||
{
|
||||
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;
|
||||
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);
|
||||
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 = (b * GCH + t) * GTC_PP + 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));
|
||||
}
|
||||
if (tid < GCH) gdn_cp16(reinterpret_cast<float*>(sc + b * GCH + tid), reinterpret_cast<const float*>(scg + ((int64_t) ch * HV + vh) * GCH + tid));
|
||||
if (tid == 0) gdn_cp4(ste + b, teg + (int64_t) ch * HV + vh);
|
||||
gdn_cp_commit();
|
||||
};
|
||||
const int c0 = 16 * warp; // the warp's value columns c0 + gid and c0 + gid + 8
|
||||
@@ -476,7 +486,7 @@ gdn_tc_scan_kernel(float* __restrict__ state, const float* __restrict__ h, const
|
||||
const float* p = state + (size_t) (8 * j + 2 * tig) * rs + (size_t) vh * S + c0 + 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);
|
||||
stage(0, 0, 3);
|
||||
const float out_scale = rsqrtf((float) S);
|
||||
for (int ch = 0; ch < nch; ++ch) {
|
||||
const int b = ch & 1;
|
||||
@@ -484,16 +494,17 @@ gdn_tc_scan_kernel(float* __restrict__ state, const float* __restrict__ h, const
|
||||
const int n = (int) min((int64_t) GCH, T - c_t0);
|
||||
gdn_cp_wait_all(); // this thread's pieces of chunk ch
|
||||
__syncthreads(); // every thread's: chunk ch's buffer complete; every read of the other one done
|
||||
if (ch + 1 < nch) stage(ch + 1, b ^ 1);
|
||||
if (ch + 1 < nch) stage(ch + 1, b ^ 1, SPLIT ? 1 : 3);
|
||||
const uint16_t* bkh = kh + b * GCH * GTC_KP;
|
||||
const uint16_t* bkl = kl + b * GCH * GTC_KP;
|
||||
const uint16_t* bqh = qh + b * GCH * GTC_KP;
|
||||
const uint16_t* bql = ql + b * GCH * GTC_KP;
|
||||
const uint16_t* bth = th + b * GCH * GTC_PP;
|
||||
const uint16_t* btl = tl + b * GCH * GTC_PP;
|
||||
const uint16_t* bph = ph + b * GCH * GTC_PP;
|
||||
const uint16_t* bpl = pl + b * GCH * GTC_PP;
|
||||
const float* bsv = sv + b * GCH * GTC_VP;
|
||||
const int bv = SPLIT ? 0 : b;
|
||||
const uint16_t* bth = th + bv * GCH * GTC_PP;
|
||||
const uint16_t* btl = tl + bv * GCH * GTC_PP;
|
||||
const uint16_t* bph = ph + bv * GCH * GTC_PP;
|
||||
const uint16_t* bpl = pl + bv * GCH * GTC_PP;
|
||||
const float* bsv = sv + bv * GCH * GTC_VP;
|
||||
const float4* bsc = sc + b * GCH;
|
||||
// (a) X^T = S0^T K^T and QS^T = S0^T Q^T, the state x 2^es (its largest entry in [2^14, 2^15)). The lane's
|
||||
// maximum over eight chains (nvcc keeps the order written: one chain was 64 dependent FMNMX before the first mma)
|
||||
@@ -558,6 +569,10 @@ gdn_tc_scan_kernel(float* __restrict__ state, const float* __restrict__ h, const
|
||||
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 (SPLIT) { // v, T and P of this chunk are dead: every warp is past (d); the next chunk's land during (e)
|
||||
__syncthreads();
|
||||
if (ch + 1 < nch) stage(ch + 1, 0, 2);
|
||||
}
|
||||
// (e) S^T = eL S0^T + Dd^T K, Dd^T = f_i D^T: the accumulator x (eL 2^(ed + 8)) before (one rounding, as the
|
||||
// s *= gL of gdn_chunk_scan_kernel), Dd^T in x 2^ed, K x 2^8, the accumulator x 2^-(ed + 8) after
|
||||
const float sPre = gtc_lds128(bsc + GCH - 1).y * sO; // exp(gL)
|
||||
|
||||
+112
-17
@@ -1,6 +1,7 @@
|
||||
// src/prefill/kernels.cu - see include/strata/prefill/kernels.hpp.
|
||||
#include "strata/prefill/kernels.hpp"
|
||||
#include "strata/prefill/gdn_chunk.hpp"
|
||||
#include "strata/prefill/gdn_probe.hpp"
|
||||
#include "strata/kernels/mrope.hpp"
|
||||
#include "strata/kernels/gfx_arch.hpp"
|
||||
#include "strata/kernels/router_top10.hpp"
|
||||
@@ -1191,13 +1192,19 @@ cudaError_t gdn_rec_chunked(float* state, const float* h, const float* gate, con
|
||||
// STRATA_GDN_CHUNKED=2 skips the SM-count test (measurement on a card under 128 SMs: the scan then runs in
|
||||
// several waves and is expected to be slower than the default kernel; the architecture and shared-memory
|
||||
// tests stay)
|
||||
static const bool any_sms = [] { const char* v = std::getenv("STRATA_GDN_CHUNKED"); return v != nullptr && std::atoi(v) >= 2; }();
|
||||
static const bool any_sms_env = [] { const char* v = std::getenv("STRATA_GDN_CHUNKED"); return v != nullptr && std::atoi(v) >= 2; }();
|
||||
const GdnProbe* pr = gdn_probe_get(); // tests only: the same as =2, per call
|
||||
const bool any_sms = any_sms_env || (pr != nullptr && pr->any_sms);
|
||||
if (strata::cc_major_of(major) < 8 || (!any_sms && sms < HK * (S / GDV)) || (size_t) strata::smem_optin_of(smem) < kChunkScanSmem)
|
||||
return cudaErrorNotSupported;
|
||||
}
|
||||
float* scratch = gdn_chunk_scratch();
|
||||
if (scratch == nullptr) return cudaErrorMemoryAllocation;
|
||||
const int nb = (int) ((std::min<int64_t>(GSB, T) + GCH - 1) / GCH);
|
||||
// the scratch segment: a test may ask for a smaller one than GSB (a multiple of GCH); production always GSB
|
||||
int64_t gsb = GSB;
|
||||
if (const GdnProbe* pr = gdn_probe_get())
|
||||
if (pr->gsb >= GCH && pr->gsb < GSB) gsb = (pr->gsb / GCH) * GCH;
|
||||
const int nb = (int) ((std::min<int64_t>(gsb, T) + GCH - 1) / GCH);
|
||||
const size_t tsz = (size_t) nb * HV * GCH * GCH;
|
||||
float *tm = scratch, *pm = tm + tsz, *gm = pm + tsz;
|
||||
{ // once a device
|
||||
@@ -1210,8 +1217,8 @@ cudaError_t gdn_rec_chunked(float* state, const float* h, const float* gate, con
|
||||
attrs[dev] = true;
|
||||
}
|
||||
}
|
||||
for (int64_t t0 = 0; t0 < T; t0 += GSB) {
|
||||
const int nch = (int) ((std::min<int64_t>(GSB, T - t0) + GCH - 1) / GCH);
|
||||
for (int64_t t0 = 0; t0 < T; t0 += gsb) {
|
||||
const int nch = (int) ((std::min<int64_t>(gsb, T - t0) + GCH - 1) / GCH);
|
||||
gdn_chunk_prep_kernel<<<dim3((unsigned) nch, HK), 256, kChunkPrepSmem, s>>>(h, gate, beta, tm, pm, gm, t0, T);
|
||||
gdn_chunk_scan_kernel<<<dim3(HK, S / GDV), 128 * VPK, kChunkScanSmem, s>>>(state, h, beta, tm, pm, gm, y, t0, nch, T);
|
||||
}
|
||||
@@ -1221,7 +1228,7 @@ cudaError_t gdn_rec_chunked(float* state, const float* h, const float* gate, con
|
||||
// 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.
|
||||
cudaError_t gdn_rec_tc(float* state, const float* h, const float* gate, const float* beta, float* y, int64_t T,
|
||||
cudaStream_t s) {
|
||||
cudaStream_t s, bool split = false) {
|
||||
if (T <= 0) return cudaSuccess;
|
||||
{
|
||||
int dev = 0, major = 0, sms = 0, smem = 0;
|
||||
@@ -1232,27 +1239,57 @@ cudaError_t gdn_rec_tc(float* state, const float* h, const float* gate, const fl
|
||||
cudaGetLastError();
|
||||
return cudaErrorNotSupported;
|
||||
}
|
||||
if (strata::cc_major_of(major) < 8 || sms < HV || (size_t) strata::smem_optin_of(smem) < kTcScanSmem<false>)
|
||||
if (strata::cc_major_of(major) < 8 || sms < HV ||
|
||||
(size_t) strata::smem_optin_of(smem) < (split ? kTcScanSmemSplit : kTcScanSmem<false>))
|
||||
return cudaErrorNotSupported;
|
||||
}
|
||||
const GtcScratch sc = gdn_tc_scratch();
|
||||
if (sc.ti == nullptr) return cudaErrorMemoryAllocation;
|
||||
{ // once a device
|
||||
static bool attrs[64] = {};
|
||||
static bool attrs[2][64] = {};
|
||||
int dev = 0;
|
||||
cudaGetDevice(&dev);
|
||||
if (dev >= 0 && dev < 64 && !attrs[dev]) {
|
||||
cudaFuncSetAttribute(gdn_tc_scan_kernel<false>, cudaFuncAttributeMaxDynamicSharedMemorySize, (int) kTcScanSmem<false>);
|
||||
attrs[dev] = true;
|
||||
if (dev >= 0 && dev < 64 && !attrs[split][dev]) {
|
||||
if (split) cudaFuncSetAttribute(gdn_tc_scan_kernel<false, true>, cudaFuncAttributeMaxDynamicSharedMemorySize, (int) kTcScanSmemSplit);
|
||||
else cudaFuncSetAttribute(gdn_tc_scan_kernel<false>, cudaFuncAttributeMaxDynamicSharedMemorySize, (int) kTcScanSmem<false>);
|
||||
attrs[split][dev] = true;
|
||||
}
|
||||
}
|
||||
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);
|
||||
gdn_tc_scan_kernel<false><<<HV, GTC_CW * 32, kTcScanSmem<false>, s>>>(state, h, sc.ti, sc.te, sc.sc, sc.kq, sc.pi, y, t0, nch, T);
|
||||
if (split)
|
||||
gdn_tc_scan_kernel<false, true><<<HV, GTC_CW * 32, kTcScanSmemSplit, s>>>(state, h, sc.ti, sc.te, sc.sc, sc.kq, sc.pi, y, t0, nch, T);
|
||||
else
|
||||
gdn_tc_scan_kernel<false><<<HV, GTC_CW * 32, kTcScanSmem<false>, s>>>(state, h, sc.ti, sc.te, sc.sc, sc.kq, sc.pi, y, t0, nch, T);
|
||||
}
|
||||
return cudaSuccess;
|
||||
}
|
||||
// Why gdn_rec_chunked (tc false) or gdn_rec_tc (tc true) launched nothing, in words: the same tests asked again. Only
|
||||
// the gdn_probe.hpp `reason` uses it (a test says what it timed); the dispatch itself never calls it.
|
||||
const char* gdn_chunk_why(bool tc, cudaError_t e, bool any_sms, bool split = false) {
|
||||
static thread_local char buf[200];
|
||||
if (e == cudaErrorMemoryAllocation) return "scratch allocation failed";
|
||||
int dev = 0, major = 0, sms = 0, smem = 0;
|
||||
if (cudaGetDevice(&dev) != cudaSuccess ||
|
||||
cudaDeviceGetAttribute(&major, cudaDevAttrComputeCapabilityMajor, dev) != cudaSuccess ||
|
||||
cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, dev) != cudaSuccess ||
|
||||
cudaDeviceGetAttribute(&smem, cudaDevAttrMaxSharedMemoryPerBlockOptin, dev) != cudaSuccess) {
|
||||
cudaGetLastError();
|
||||
return "device attributes unavailable";
|
||||
}
|
||||
const int cc = strata::cc_major_of(major);
|
||||
const size_t need = tc ? (split ? kTcScanSmemSplit : (size_t) kTcScanSmem<false>) : kChunkScanSmem;
|
||||
const int sm_need = tc ? HV : HK * (S / GDV);
|
||||
if (cc < 8) std::snprintf(buf, sizeof buf, "compute capability %d < 8 (cp.async)", cc);
|
||||
else if (!tc && !any_sms && sms < sm_need)
|
||||
std::snprintf(buf, sizeof buf, "%d SMs < %d: the scan would run in several waves (STRATA_GDN_CHUNKED=2 bypasses)", sms, sm_need);
|
||||
else if (tc && sms < sm_need) std::snprintf(buf, sizeof buf, "%d SMs < %d (one block per value head)", sms, sm_need);
|
||||
else if ((size_t) strata::smem_optin_of(smem) < need)
|
||||
std::snprintf(buf, sizeof buf, "shared memory per block %d B < %zu B needed", strata::smem_optin_of(smem), need);
|
||||
else std::snprintf(buf, sizeof buf, "not supported (%s)", cudaGetErrorString(e));
|
||||
return buf;
|
||||
}
|
||||
#endif
|
||||
#if defined(__HIPCC__) // the Aurora (gfx11) recurrence kernels: AMD builds only
|
||||
// Aurora (S23): the recurrence with the column-split grid of gdn_rec_cols_pipe_kernel (four 32-column blocks per value
|
||||
@@ -2777,16 +2814,74 @@ static void gdn_recurrence_impl(int variant, float* state, const float* h, const
|
||||
// chunks of kGdnTcMin+ tokens: within 1e-4 of the kernels below (gdn_rec_parity), 3.75x gdn_rec_kh_kernel on a
|
||||
// CMP 170HX; a shorter chunk, a card it does not fit or a failed scratch allocation takes the kernels below
|
||||
static const bool tc = [] { const char* v = std::getenv("STRATA_GDN_TC"); return v != nullptr && std::atoi(v) != 0; }();
|
||||
if (tc && T >= kGdnTcMin && gdn_rec_tc(state, h, gate, beta, y, T, (cudaStream_t) stream) == cudaSuccess) {
|
||||
} else if (chunked && T >= kGdnChunkedMin && gdn_rec_chunked(state, h, gate, beta, y, T, (cudaStream_t) stream) == cudaSuccess) {
|
||||
} else if (pipe && gdn_keyhead_ok()) // the value heads of a key head in one thread (same bits)
|
||||
gdn_rec_kh_kernel<<<HK * NCB, dim3(CB, RG), 0, (cudaStream_t) stream>>>(state, h, gate, beta, y, T);
|
||||
else
|
||||
#endif
|
||||
// A test's probe (gdn_probe.hpp) reads back the path that ran and may force one; without a probe this is the
|
||||
// dispatch it was. A forced path that cannot run launches nothing (selected = none, reason says why).
|
||||
GdnProbe* pr = gdn_probe_get();
|
||||
const int force = pr ? pr->force : (int) kGdnAuto;
|
||||
int sel = kGdnNone, wanted = force;
|
||||
const char* why = "";
|
||||
const cudaStream_t cs = (cudaStream_t) stream;
|
||||
if (force != kGdnAuto) {
|
||||
if (force == kGdnTc || force == kGdnTcSplit) {
|
||||
const bool sp = force == kGdnTcSplit;
|
||||
cudaError_t e = T >= kGdnTcMin ? gdn_rec_tc(state, h, gate, beta, y, T, cs, sp) : cudaErrorInvalidValue;
|
||||
if (e == cudaSuccess) sel = force;
|
||||
else why = T < kGdnTcMin ? "T below the tc minimum (128)" : gdn_chunk_why(true, e, false, sp);
|
||||
} else if (force == kGdnChunked) {
|
||||
cudaError_t e = T >= kGdnChunkedMin ? gdn_rec_chunked(state, h, gate, beta, y, T, cs) : cudaErrorInvalidValue;
|
||||
if (e == cudaSuccess) sel = kGdnChunked;
|
||||
else why = T < kGdnChunkedMin ? "T below the chunked minimum (128)" : gdn_chunk_why(false, e, pr->any_sms);
|
||||
} else if (force == kGdnKeyHead) {
|
||||
if (gdn_keyhead_ok()) {
|
||||
gdn_rec_kh_kernel<<<HK * NCB, dim3(CB, RG), 0, cs>>>(state, h, gate, beta, y, T);
|
||||
sel = kGdnKeyHead;
|
||||
} else why = "keyhead kernel not eligible here (needs sm_80+ with all 64 blocks resident, or STRATA_GDN_KEYHEAD=0 is set)";
|
||||
} else if (force == kGdnPipe) {
|
||||
gdn_rec_cols_pipe_kernel<<<HV * NCB, dim3(CB, RG), 0, cs>>>(state, h, gate, beta, y, T);
|
||||
sel = kGdnPipe;
|
||||
} else if (force == kGdnCols) {
|
||||
gdn_rec_cols_kernel<<<HV * NCB, dim3(CB, RG), 0, cs>>>(state, h, gate, beta, y, T);
|
||||
sel = kGdnCols;
|
||||
} else why = "unknown path";
|
||||
} else {
|
||||
if (tc && T >= kGdnTcMin) {
|
||||
wanted = kGdnTc;
|
||||
cudaError_t e = gdn_rec_tc(state, h, gate, beta, y, T, cs);
|
||||
if (e == cudaSuccess) sel = kGdnTc;
|
||||
else { // the 99-KB cards (sm_86 / 89 / 120): the single-buffered v / T / P plan
|
||||
if (pr) why = gdn_chunk_why(true, e, false);
|
||||
e = gdn_rec_tc(state, h, gate, beta, y, T, cs, true);
|
||||
if (e == cudaSuccess) { sel = kGdnTcSplit; why = ""; }
|
||||
}
|
||||
}
|
||||
if (sel == kGdnNone && chunked && T >= kGdnChunkedMin) {
|
||||
if (wanted == kGdnAuto) wanted = kGdnChunked;
|
||||
const cudaError_t e = gdn_rec_chunked(state, h, gate, beta, y, T, cs);
|
||||
if (e == cudaSuccess) sel = kGdnChunked;
|
||||
else if (pr && wanted == kGdnChunked) why = gdn_chunk_why(false, e, false);
|
||||
}
|
||||
if (sel == kGdnNone) {
|
||||
if (pipe && gdn_keyhead_ok()) { // the value heads of a key head in one thread (same bits)
|
||||
gdn_rec_kh_kernel<<<HK * NCB, dim3(CB, RG), 0, cs>>>(state, h, gate, beta, y, T);
|
||||
sel = kGdnKeyHead;
|
||||
} else if (pipe) { // the software-pipelined loads (same bits)
|
||||
gdn_rec_cols_pipe_kernel<<<HV * NCB, dim3(CB, RG), 0, cs>>>(state, h, gate, beta, y, T);
|
||||
sel = kGdnPipe;
|
||||
} else {
|
||||
gdn_rec_cols_kernel<<<HV * NCB, dim3(CB, RG), 0, cs>>>(state, h, gate, beta, y, T);
|
||||
sel = kGdnCols;
|
||||
}
|
||||
if (wanted == kGdnAuto) wanted = sel;
|
||||
}
|
||||
}
|
||||
if (pr) { pr->selected = sel; pr->wanted = wanted == kGdnAuto ? sel : wanted; pr->reason = why; ++pr->calls; }
|
||||
if (sel == kGdnNone) { check("gdn_recurrence (forced path not eligible: nothing launched)"); return; }
|
||||
#else
|
||||
if (pipe) // the software-pipelined loads (same bits)
|
||||
gdn_rec_cols_pipe_kernel<<<HV * NCB, dim3(CB, RG), 0, (cudaStream_t) stream>>>(state, h, gate, beta, y, T);
|
||||
else
|
||||
gdn_rec_cols_kernel<<<HV * NCB, dim3(CB, RG), 0, (cudaStream_t) stream>>>(state, h, gate, beta, y, T);
|
||||
#endif
|
||||
if (with_out)
|
||||
gdn_out_norm_kernel<<<dim3((unsigned) T, HV), S, 0, (cudaStream_t) stream>>>(z, gamma, eps, y, y16, ld16);
|
||||
}
|
||||
|
||||
@@ -0,0 +1,245 @@
|
||||
// tests/gdn/gdn_fixture.hpp - backend-independent fixtures and an FP64 oracle for the Flash-Next GDN prompt path.
|
||||
//
|
||||
// Plain C++17, no GPU headers: the CUDA harness (src/prefill/gdn_harness.cu) uses it today; the HIP and Metal/SYCL
|
||||
// harnesses are meant to include the same file, so every backend is judged against the same inputs and the same
|
||||
// reference. Contract: docs/GDN_CONTRACT.md (shapes, layouts, what a candidate may change).
|
||||
//
|
||||
// Fixture raw inputs of one layer for T tokens: qkv [T][C], ab [T][2 HV], dt, ssm_a, conv_w [C][4], conv
|
||||
// history [C][3], the recurrent state [S][HV][S] (row, head, column), z [T][HV][S], gamma [S]
|
||||
// RefResult FP64 oracle for some value heads over the whole sequence: normalised output (what the out
|
||||
// projection reads, before the FP16 rounding), the final state of those heads, the conv history
|
||||
// ErrStat max_abs, normalized_max (max_abs / max|ref|), relative_L2 (gdn-research 06)
|
||||
// The oracle starts from the RAW inputs (conv, SiLU, L2 norm, gates, recurrence, output norm), so a bug shared by a
|
||||
// kernel and a test copy of it cannot hide.
|
||||
#pragma once
|
||||
|
||||
#include <algorithm>
|
||||
#include <cmath>
|
||||
#include <cstdint>
|
||||
#include <cstdio>
|
||||
#include <cstring>
|
||||
#include <thread>
|
||||
#include <vector>
|
||||
|
||||
namespace gdnfx {
|
||||
|
||||
constexpr int S = 128, HK = 16, HV = 48, VPK = HV / HK, C = 2 * HK * S + HV * S; // C = 10240
|
||||
constexpr float EPS = 1e-6f; // the engine's L2 / RMS epsilon
|
||||
|
||||
struct Rng { // splitmix64: identical on every platform and compiler
|
||||
uint64_t s;
|
||||
explicit Rng(uint64_t seed) : s(seed * 0x9E3779B97F4A7C15ULL + 0x1234567ULL) {}
|
||||
uint64_t next() {
|
||||
uint64_t z = (s += 0x9E3779B97F4A7C15ULL);
|
||||
z = (z ^ (z >> 30)) * 0xBF58476D1CE4E5B9ULL;
|
||||
z = (z ^ (z >> 27)) * 0x94D049BB133111EBULL;
|
||||
return z ^ (z >> 31);
|
||||
}
|
||||
double uni() { return (double) (next() >> 11) * (1.0 / 9007199254740992.0); } // [0, 1)
|
||||
double gauss() {
|
||||
double u1 = uni(), u2 = uni();
|
||||
if (u1 < 1e-300) u1 = 1e-300;
|
||||
return std::sqrt(-2.0 * std::log(u1)) * std::cos(6.283185307179586 * u2);
|
||||
}
|
||||
};
|
||||
|
||||
struct FixtureSpec {
|
||||
int64_t T = 0;
|
||||
uint64_t seed = 1;
|
||||
// state0: "random" (as after an earlier prompt chunk), "zero" (a fresh conversation), "big", "small"
|
||||
const char* state = "random";
|
||||
float v_scale = 1.0f; // values x this (50: range stress)
|
||||
float gate_scale = 1.0f; // ssm_a x this (0.005: near-zero decay, 3: strong decay)
|
||||
bool coherent_keys = false; // every token's k is a head constant plus a small variation (T = (I+A)^-1 large)
|
||||
bool alt_beta = false; // beta alternates 0.02 / 0.98 per token (endpoint regime)
|
||||
bool tiny_k = false; // every 7th token's k/q channels are ~1e-4 before the L2 norm (eps matters)
|
||||
};
|
||||
|
||||
struct Fixture {
|
||||
FixtureSpec spec;
|
||||
int64_t T = 0;
|
||||
std::vector<float> qkv, ab, dt, ssm_a, conv_w, hist, state0, z, gamma;
|
||||
};
|
||||
|
||||
inline Fixture make_fixture(const FixtureSpec& sp) {
|
||||
Fixture f;
|
||||
f.spec = sp;
|
||||
f.T = sp.T;
|
||||
const int64_t T = sp.T;
|
||||
Rng r(sp.seed);
|
||||
f.qkv.resize((size_t) T * C);
|
||||
f.ab.resize((size_t) T * 2 * HV);
|
||||
f.dt.resize(HV);
|
||||
f.ssm_a.resize(HV);
|
||||
f.conv_w.resize((size_t) C * 4);
|
||||
f.hist.resize((size_t) C * 3);
|
||||
f.state0.resize((size_t) S * HV * S);
|
||||
f.z.resize((size_t) T * HV * S);
|
||||
f.gamma.resize(S);
|
||||
for (auto& x : f.qkv) x = (float) r.gauss();
|
||||
if (sp.coherent_keys) {
|
||||
std::vector<float> base((size_t) HK * S);
|
||||
for (auto& x : base) x = (float) r.gauss();
|
||||
for (int64_t t = 0; t < T; ++t)
|
||||
for (int i = 0; i < HK * S; ++i) {
|
||||
float& x = f.qkv[(size_t) t * C + HK * S + i]; // k channels
|
||||
x = base[i] + 0.05f * x;
|
||||
}
|
||||
}
|
||||
if (sp.tiny_k)
|
||||
for (int64_t t = 0; t < T; t += 7)
|
||||
for (int i = 0; i < 2 * HK * S; ++i) f.qkv[(size_t) t * C + i] *= 1e-4f;
|
||||
if (sp.v_scale != 1.0f)
|
||||
for (int64_t t = 0; t < T; ++t)
|
||||
for (int i = 0; i < HV * S; ++i) f.qkv[(size_t) t * C + 2 * HK * S + i] *= sp.v_scale;
|
||||
for (auto& x : f.ab) x = (float) r.gauss();
|
||||
if (sp.alt_beta)
|
||||
for (int64_t t = 0; t < T; ++t)
|
||||
for (int h = 0; h < HV; ++h) f.ab[(size_t) t * 2 * HV + HV + h] = (t & 1) ? 3.9f : -3.9f;
|
||||
for (int h = 0; h < HV; ++h) {
|
||||
f.dt[h] = (float) (0.5 * r.gauss() - 0.5);
|
||||
const double u = r.uni();
|
||||
f.ssm_a[h] = (float) (-(0.05 + 1.5 * u * u) * sp.gate_scale);
|
||||
}
|
||||
for (auto& x : f.conv_w) x = (float) (0.35 * r.gauss());
|
||||
for (auto& x : f.hist) x = (float) r.gauss();
|
||||
for (auto& x : f.z) x = (float) r.gauss();
|
||||
for (auto& x : f.gamma) x = (float) (1.0 + 0.1 * r.gauss());
|
||||
float ss = 0.05f;
|
||||
if (!std::strcmp(sp.state, "zero")) ss = 0.0f;
|
||||
else if (!std::strcmp(sp.state, "big")) ss = 2.5f;
|
||||
else if (!std::strcmp(sp.state, "small")) ss = 5e-4f;
|
||||
for (auto& x : f.state0) x = ss == 0.0f ? 0.0f : (float) (ss * r.gauss());
|
||||
return f;
|
||||
}
|
||||
|
||||
// ------------------------------------------------------------------------------------------------ FP64 oracle
|
||||
inline double sigm(double x) { return 1.0 / (1.0 + std::exp(-x)); }
|
||||
|
||||
struct RefResult {
|
||||
std::vector<int> heads; // value heads computed
|
||||
std::vector<std::vector<double>> y; // per head: [T][S] normalised output (before FP16)
|
||||
std::vector<std::vector<double>> state; // per head: final [S][S] (row, column)
|
||||
std::vector<float> hist; // conv history after the whole sequence (exact: copies of inputs)
|
||||
};
|
||||
|
||||
// One value head over tokens [0, T). state_io: that head's [S][S] (row, column), updated in place.
|
||||
inline void ref_head(const Fixture& f, int vh, std::vector<double>& st, std::vector<double>& y) {
|
||||
const int64_t T = f.T;
|
||||
const int qh = vh % HK;
|
||||
y.assign((size_t) T * S, 0.0);
|
||||
std::vector<double> q(S), k(S), v(S), kv(S), delta(S), o(S);
|
||||
auto in = [&](int64_t t, int c) -> double {
|
||||
return t >= 0 ? (double) f.qkv[(size_t) t * C + c] : (double) f.hist[(size_t) c * 3 + (int) (t + 3)];
|
||||
};
|
||||
auto conv = [&](int64_t t, int c) -> double {
|
||||
const double s = (double) f.conv_w[(size_t) c * 4] * in(t - 3, c) + (double) f.conv_w[(size_t) c * 4 + 1] * in(t - 2, c) +
|
||||
(double) f.conv_w[(size_t) c * 4 + 2] * in(t - 1, c) + (double) f.conv_w[(size_t) c * 4 + 3] * in(t, c);
|
||||
return s / (1.0 + std::exp(-s));
|
||||
};
|
||||
for (int64_t t = 0; t < T; ++t) {
|
||||
double qq = 0, kk = 0;
|
||||
for (int i = 0; i < S; ++i) {
|
||||
q[i] = conv(t, qh * S + i);
|
||||
k[i] = conv(t, HK * S + qh * S + i);
|
||||
v[i] = conv(t, 2 * HK * S + vh * S + i);
|
||||
qq += q[i] * q[i];
|
||||
kk += k[i] * k[i];
|
||||
}
|
||||
const double qs = 1.0 / std::sqrt(qq + (double) EPS), ks = 1.0 / std::sqrt(kk + (double) EPS);
|
||||
for (int i = 0; i < S; ++i) { q[i] *= qs; k[i] *= ks; }
|
||||
const double a = (double) f.ab[(size_t) t * 2 * HV + vh] + (double) f.dt[vh];
|
||||
const double sp = a > 20.0 ? a : std::log1p(std::exp(a));
|
||||
const double g = std::exp(sp * (double) f.ssm_a[vh]);
|
||||
const double beta = sigm((double) f.ab[(size_t) t * 2 * HV + HV + vh]);
|
||||
for (int c = 0; c < S; ++c) kv[c] = 0;
|
||||
for (int i = 0; i < S; ++i) {
|
||||
const double* row = &st[(size_t) i * S];
|
||||
for (int c = 0; c < S; ++c) kv[c] += row[c] * k[i];
|
||||
}
|
||||
for (int c = 0; c < S; ++c) delta[c] = (v[c] - g * kv[c]) * beta;
|
||||
for (int c = 0; c < S; ++c) o[c] = 0;
|
||||
for (int i = 0; i < S; ++i) {
|
||||
double* row = &st[(size_t) i * S];
|
||||
for (int c = 0; c < S; ++c) {
|
||||
row[c] = g * row[c] + k[i] * delta[c];
|
||||
o[c] += row[c] * q[i];
|
||||
}
|
||||
}
|
||||
double ss = 0;
|
||||
for (int c = 0; c < S; ++c) { o[c] /= std::sqrt((double) S); ss += o[c] * o[c]; }
|
||||
const double rn = 1.0 / std::sqrt(ss / S + (double) EPS);
|
||||
for (int c = 0; c < S; ++c)
|
||||
y[(size_t) t * S + c] = o[c] * rn * (double) f.gamma[c] * sigm((double) f.z[((size_t) t * HV + vh) * S + c]);
|
||||
}
|
||||
}
|
||||
|
||||
// The oracle for `heads`, one thread each (state taken from the fixture's state0, layout [row][head][col]).
|
||||
inline RefResult ref_run(const Fixture& f, const std::vector<int>& heads) {
|
||||
RefResult r;
|
||||
r.heads = heads;
|
||||
r.y.resize(heads.size());
|
||||
r.state.resize(heads.size());
|
||||
std::vector<std::thread> th;
|
||||
for (size_t n = 0; n < heads.size(); ++n)
|
||||
th.emplace_back([&, n] {
|
||||
const int vh = heads[n];
|
||||
std::vector<double> st((size_t) S * S);
|
||||
for (int i = 0; i < S; ++i)
|
||||
for (int c = 0; c < S; ++c) st[(size_t) i * S + c] = (double) f.state0[((size_t) i * HV + vh) * S + c];
|
||||
ref_head(f, vh, st, r.y[n]);
|
||||
r.state[n] = std::move(st);
|
||||
});
|
||||
for (auto& t : th) t.join();
|
||||
r.hist.resize((size_t) C * 3);
|
||||
for (int c = 0; c < C; ++c)
|
||||
for (int k = 0; k < 3; ++k) {
|
||||
const int64_t t = f.T - 3 + k;
|
||||
r.hist[(size_t) c * 3 + k] = t >= 0 ? f.qkv[(size_t) t * C + c] : f.hist[(size_t) c * 3 + (int) (t + 3)];
|
||||
}
|
||||
return r;
|
||||
}
|
||||
|
||||
// ------------------------------------------------------------------------------------------------ metrics
|
||||
struct ErrStat {
|
||||
double max_abs = 0, ref_max = 0, sq_diff = 0, sq_ref = 0;
|
||||
long nonfinite = 0, n = 0;
|
||||
void add(double got, double ref) {
|
||||
++n;
|
||||
if (!std::isfinite(got)) { ++nonfinite; return; }
|
||||
const double d = std::fabs(got - ref);
|
||||
if (d > max_abs) max_abs = d;
|
||||
const double a = std::fabs(ref);
|
||||
if (a > ref_max) ref_max = a;
|
||||
sq_diff += d * d;
|
||||
sq_ref += ref * ref;
|
||||
}
|
||||
double normalized_max() const { return max_abs / std::max(ref_max, 1e-30); }
|
||||
double rel_l2() const { return std::sqrt(sq_diff) / std::max(std::sqrt(sq_ref), 1e-30); }
|
||||
};
|
||||
|
||||
inline uint64_t fnv1a(const void* p, size_t bytes) {
|
||||
const unsigned char* b = (const unsigned char*) p;
|
||||
uint64_t h = 1469598103934665603ULL;
|
||||
for (size_t i = 0; i < bytes; ++i) { h ^= b[i]; h *= 1099511628211ULL; }
|
||||
return h;
|
||||
}
|
||||
|
||||
inline float half_to_float(uint16_t h) { // IEEE binary16 -> float, host side
|
||||
const uint32_t sign = (uint32_t) (h >> 15) << 31;
|
||||
uint32_t e = (h >> 10) & 0x1f, m = h & 0x3ff, out;
|
||||
if (e == 0) {
|
||||
if (m == 0) out = sign;
|
||||
else {
|
||||
e = 127 - 15 + 1;
|
||||
while (!(m & 0x400)) { m <<= 1; --e; }
|
||||
out = sign | (e << 23) | ((m & 0x3ff) << 13);
|
||||
}
|
||||
} else if (e == 31) out = sign | 0x7f800000u | (m << 13);
|
||||
else out = sign | ((e + 112) << 23) | (m << 13);
|
||||
float f;
|
||||
std::memcpy(&f, &out, 4);
|
||||
return f;
|
||||
}
|
||||
|
||||
} // namespace gdnfx
|
||||
Reference in new issue
Block a user