mirror of
https://github.com/Niko1221/Strata.git
synced 2026-11-07 18:35:04 -10:00
verify: fewer launches around the shared expert and the KV append, bitwise (opt-in STRATA_LFUSE=1)
Per MoE layer of a verify window (2-8 rows, resident combine path): - the shared expert's scalar gate (a 1-row BF16 GEMV, one latency-bound workgroup) rides as one more block of the router GEMV launch (bf16_gemv_fp32_mmvf_multi_aux; the rows kernel's per-output code is the multi kernel's - checked bitwise for 2-8 rows, with and without TSUM) - its sigmoid + row scale move into the combine (combine<true, true>: the same sigmoid expression, the product rounded on its own with contraction off in that block only, then added as before) - shared gate and up projections in one launch (native_mmvq_pair: blocks n_out.. take the second matrix; Q8_0, exact multi layout, not packed) Per attention layer: the window's int8 KV appends in one launch (kv_append_q8_steps, blockIdx.z = token x K/V; distinct cells). -4 launches per MoE layer, -(n-1) per attention layer; STRATA_LFUSE_GATE=0 / STRATA_LFUSE_PAIR=0 turn the gate / pair parts off. Engine (gfx1151, UD-IQ4_XS, c-best3 + EXPERT_V2 + Q8_PACKED + TSUM, S25's 18 prompts, same binary): 51.29 -> 50.31 ms per verify window, 54.33 -> 55.12 tok/s; 59K prompts 54.19 -> 53.17 ms, 42.97 -> 43.79 tok/s; all outputs identical.
This commit is contained in:
1 parent
75c9fc6baa
commit
b08cf3e160
11 files changed
+204
-17
No files matched your search
@@ -48,5 +48,11 @@ void bf16_gemv_fp32_mmvf(const float* x, const uint16_t* w, float* y,
|
||||
/// the weight read once, every output bit-identical to its own single-row call.
|
||||
void bf16_gemv_fp32_mmvf_multi(const float* x, int64_t ldx, const uint16_t* w, float* y, int64_t ldy,
|
||||
int64_t n_in, int64_t n_out, int n_tok, void* stream);
|
||||
/// S26 STRATA_LFUSE: bf16_gemv_fp32_mmvf_multi(x, ldx, w, y, ldy, n_in, n_out, n_tok) plus the 1-row
|
||||
/// bf16_gemv_fp32_mmvf_multi(x, ldx, w_aux, y_aux, ldy_aux, n_in, 1, n_tok) in ONE launch, every output bitwise the
|
||||
/// two calls'. Returns false (nothing launched) unless the rows path applies (STRATA_MMVF_ROWS=1, n_out >= 64, 2-8 rows).
|
||||
bool bf16_gemv_fp32_mmvf_multi_aux(const float* x, int64_t ldx, const uint16_t* w, float* y, int64_t ldy,
|
||||
int64_t n_in, int64_t n_out, int n_tok, const uint16_t* w_aux, float* y_aux,
|
||||
int64_t ldy_aux, void* stream);
|
||||
|
||||
} // namespace strata::kernels
|
||||
@@ -28,6 +28,11 @@ inline uint64_t kv_q8_bytes_per_cell(const QsaShapes& s) {
|
||||
|
||||
/// Append the cell at step[kStepPos] (graph-capturable: position and page come from device memory). With a host
|
||||
/// copy (KV streaming, `kv_stream.hpp`) the cell is written there too, and to VRAM only if its block is resident.
|
||||
/// S26 STRATA_LFUSE: n_tok kv_append_q8_step calls (token j: step + j * step_stride, kcur / vcur + j * cur_stride)
|
||||
/// in one launch, bitwise the same stores (the tokens' cells are distinct).
|
||||
void kv_append_q8_steps(int8_t* k_q, int8_t* v_q, uint16_t* k_scale, uint16_t* v_scale, const int32_t* page_table,
|
||||
const int32_t* step, int step_stride, const float* kcur, const float* vcur, int cur_stride,
|
||||
int n_tok, const QsaShapes& s, void* stream, const KvHostPools* host);
|
||||
void kv_append_q8_step(int8_t* k_q, int8_t* v_q, uint16_t* k_scale, uint16_t* v_scale, const int32_t* page_table,
|
||||
const int32_t* step, const float* kcur, const float* vcur, const QsaShapes& s, void* stream,
|
||||
const KvHostPools* host = nullptr);
|
||||
|
||||
@@ -126,5 +126,9 @@ bool native_mmvq_supported(int ggml_type) noexcept;
|
||||
std::size_t native_mmvq_weight_bytes(int ggml_type, int n_in, int n_out);
|
||||
void native_mmvq(int ggml_type, const void* weights, const void* x_q8_1, float* y,
|
||||
int n_in, int n_out, int ncols, void* stream);
|
||||
/// S26 STRATA_LFUSE: native_mmvq(w1 -> y1) and native_mmvq(w2 -> y2), same type / shape / input, in ONE launch,
|
||||
/// bitwise the two calls. Returns false (nothing launched) where that is not the case (Q8_0, 2-8 columns only).
|
||||
bool native_mmvq_pair(int ggml_type, const void* w1, const void* w2, const void* x_q8_1, float* y1, float* y2,
|
||||
int n_in, int n_out, int ncols, void* stream);
|
||||
|
||||
} // namespace strata::kernels
|
||||
@@ -22,4 +22,9 @@ void native_moe_combine_multi(const float* parts, const float* weights, const fl
|
||||
/// moe_hit_add (parts += hit) gave, without the zero fill and the add launch.
|
||||
void native_moe_combine_multi_hits(const float* hits, const float* weights, const float* shared, float* output,
|
||||
int64_t n_embd, int64_t k, int n_tok, void* stream);
|
||||
/// S26 STRATA_LFUSE: the same with `shared` unscaled and shared_gate[t] the shared expert's raw gate logit: the
|
||||
/// row is scaled by sigmoid(shared_gate[t]) here, bitwise what shared_expert_multi's sigmoid + scale gave.
|
||||
void native_moe_combine_multi_hits_gated(const float* hits, const float* weights, const float* shared,
|
||||
const float* shared_gate, float* output, int64_t n_embd, int64_t k, int n_tok,
|
||||
void* stream);
|
||||
}
|
||||
@@ -78,7 +78,8 @@ void shared_expert(const uint8_t* x_q8_0, const uint8_t* x_q8k, const uint16_t*
|
||||
/// its q8_1 run as one launch (same values). Null: the separate launches.
|
||||
void shared_expert_multi(int n_tok, const float* x, const uint16_t* x_bf16, const NativeSharedWeights& nw,
|
||||
const uint16_t* gate_inp_bf16, float* gate, float* up, float* g, float* out, int64_t n_embd,
|
||||
int64_t n_ff, void* stream, const void* x_q8_1_ready = nullptr);
|
||||
int64_t n_ff, void* stream, const void* x_q8_1_ready = nullptr,
|
||||
int lfuse = 0); // S26 STRATA_LFUSE: bit 0 gate computed + applied by the caller, bit 1 gate/up pair
|
||||
|
||||
/// The MoE block's final combination, `ref/moe.py::moe` L156:
|
||||
///
|
||||
|
||||
+31
-1
@@ -56,6 +56,9 @@ constexpr float EPS = 1e-6f;
|
||||
using Clock = std::chrono::steady_clock;
|
||||
double ms_since(Clock::time_point t) { return std::chrono::duration<double, std::milli>(Clock::now() - t).count(); }
|
||||
const bool g_dbg = std::getenv("STRATA_VERIFY_DEBUG") != nullptr;
|
||||
const bool g_lfuse = [] { const char* v = std::getenv("STRATA_LFUSE"); return v != nullptr && v[0] == '1'; }();
|
||||
const bool g_lfuse_gate = [] { const char* v = std::getenv("STRATA_LFUSE_GATE"); return v == nullptr || v[0] != '0'; }();
|
||||
const bool g_lfuse_pair = [] { const char* v = std::getenv("STRATA_LFUSE_PAIR"); return v == nullptr || v[0] != '0'; }();
|
||||
const bool g_qdedup = [] { const char* v = std::getenv("STRATA_VERIFY_QDEDUP"); return v != nullptr && std::atoi(v) != 0; }();
|
||||
// S26 STRATA_QFUSE=1: activation q8_1 images written by their producers (the GDN output norm) - the same bytes
|
||||
const bool g_qfuse = [] { const char* v = std::getenv("STRATA_QFUSE"); return v != nullptr && std::atoi(v) != 0; }();
|
||||
@@ -436,6 +439,12 @@ bool Verifier::record_window(int T, cudaStream_t cs, std::string& err) {
|
||||
static const bool dec_batch = [] { const char* v = std::getenv("STRATA_DEC_BATCH"); return v == nullptr || std::atoi(v) != 0; }();
|
||||
auto stamp = [&](int64_t l, int i, int grp) { if (prof_on_ && grp == 0) gpu_stamp(prof_, (int) (l * kProfPer + i), cs); };
|
||||
const int tb_[2] = {0, (T + 1) / 2}, te_[2] = {G == 2 ? (T + 1) / 2 : T, T};
|
||||
// S26 STRATA_LFUSE=1: where the shared expert's gate / scale fusions apply (the paths they replace are the ones taken)
|
||||
auto lfuse_on = [&](int n) {
|
||||
return g_lfuse && resident_ && native_moe_combine_enabled() && dec_batch && n > 1 && n <= 8 &&
|
||||
shared_expert_native_bf16();
|
||||
};
|
||||
bool sg_gated_[2] = {false, false}; // per group: the combine applies the shared gate
|
||||
groups_[T] = G;
|
||||
|
||||
// ---- the window's inputs, from mapped staging
|
||||
@@ -644,6 +653,12 @@ bool Verifier::record_window(int T, cudaStream_t cs, std::string& err) {
|
||||
}
|
||||
stamp(l, 8, grp);
|
||||
if (grp == 0) copy_from_mapped(tail_snap_ + (size_t) qi * TS, st.idx_tail, TS, cs);
|
||||
const bool kv_multi = g_lfuse && !st.kv_hybrid && !st.kv_q4 && st.kv_int8 && n > 1; // S26 STRATA_LFUSE
|
||||
if (kv_multi)
|
||||
kv_append_q8_steps(st.k_q, st.v_q, st.k_scale, st.v_scale, st.page_table, step_ + tb * kStepCount,
|
||||
(int) kStepCount, kcur_ + tb * NKV * HD, vcur_ + tb * NKV * HD, (int) (NKV * HD), n,
|
||||
s, cs, &st.host);
|
||||
else
|
||||
for (int t = tb; t < te; ++t) {
|
||||
const int32_t* step_t = step_ + t * kStepCount;
|
||||
if (st.kv_hybrid) { // K8V4: the unused half's lanes folded onto the used pool (layer.cpp)
|
||||
@@ -734,8 +749,17 @@ bool Verifier::record_window(int T, cudaStream_t cs, std::string& err) {
|
||||
// the window's rows routed in 2 launches (one router GEMV reading the weight once, one
|
||||
// top-10) instead of 2 per token; every row's arithmetic is the single-token call's (STRATA_DEC_BATCH=0: old)
|
||||
const WeightRef* w_router = v.get("ffn_gate_inp.weight");
|
||||
// S26 STRATA_LFUSE=1: the shared expert's gate row rides in the router GEMV launch, its sigmoid + row scale
|
||||
// move into the combine, gate + up share one launch (all bitwise; resident combine path, 2-8 rows only)
|
||||
const WeightRef* w_sgi = lfuse_on(n) && g_lfuse_gate ? v.get("ffn_gate_inp_shexp.weight") : nullptr;
|
||||
bool sg_ready = false;
|
||||
if (dec_batch && n > 1 && w_router != nullptr && native_router_enabled() && NE == 512 && K == 10) {
|
||||
try {
|
||||
if (w_sgi != nullptr)
|
||||
sg_ready = bf16_gemv_fp32_mmvf_multi_aux(mixed_ + tb * N, N, (const uint16_t*) w_router->data,
|
||||
logits_ + tb * NE, NE, N, NE, n,
|
||||
(const uint16_t*) w_sgi->data, sh_g_ + tb, 1, cs);
|
||||
if (!sg_ready)
|
||||
bf16_gemv_fp32_mmvf_multi(mixed_ + tb * N, N, (const uint16_t*) w_router->data, logits_ + tb * NE, NE, N,
|
||||
NE, n, cs);
|
||||
native_router_top10_multi(logits_ + tb * NE, ids_ + tb * K, w_ + tb * K, n, cs);
|
||||
@@ -779,7 +803,9 @@ bool Verifier::record_window(int T, cudaStream_t cs, std::string& err) {
|
||||
try {
|
||||
shared_expert_multi(n, xm, sh_bf16_ + tb * N, nsw, (const uint16_t*) wgi->data, sh_gate_ + (size_t) tb * g.n_ff,
|
||||
sh_up_ + (size_t) tb * g.n_ff, sh_g_ + tb, shared_ + tb * N, N, g.n_ff, cs,
|
||||
qdedup ? (const void*) (nat_xq_ + (size_t) tb * (N / 32) * 36) : nullptr);
|
||||
qdedup ? (const void*) (nat_xq_ + (size_t) tb * (N / 32) * 36) : nullptr,
|
||||
(sg_ready ? 1 : 0) | (lfuse_on(n) && g_lfuse_pair ? 2 : 0));
|
||||
sg_gated_[grp] = sg_ready;
|
||||
} catch (const std::exception& e) {
|
||||
err = std::string("verify shared expert: ") + e.what();
|
||||
return false;
|
||||
@@ -838,6 +864,10 @@ bool Verifier::record_window(int T, cudaStream_t cs, std::string& err) {
|
||||
if (resident_ && native_moe_combine_enabled()) {
|
||||
// every row is a hit: combine straight from hit_out as 0.0f + hit (= the zeroed row + moe_hit_add)
|
||||
try {
|
||||
if (sg_gated_[grp])
|
||||
native_moe_combine_multi_hits_gated(hit_out, w_ + tb * K, shared_ + tb * N, sh_g_ + tb, bo_ + tb * N,
|
||||
N, K, n, cs);
|
||||
else
|
||||
native_moe_combine_multi_hits(hit_out, w_ + tb * K, shared_ + tb * N, bo_ + tb * N, N, K, n, cs);
|
||||
} catch (const std::exception& e) { err = "verify combine: " + std::string(e.what()); return false; }
|
||||
} else {
|
||||
|
||||
@@ -27,14 +27,22 @@ void validate(const QsaShapes& s, const char* what) {
|
||||
}
|
||||
|
||||
// One block = one 64-value group of one KV head of K (blockIdx.z = 0) or V (1); 64 threads, one value each.
|
||||
// MULTI (S26 STRATA_LFUSE): blockIdx.z = 2 * token + (K / V); token j reads step + j * step_stride and its rows at
|
||||
// kcur / vcur + j * cur_stride - every token's code is the single launch's
|
||||
template <bool MULTI = false>
|
||||
__global__ void kv_append_q8_kernel(int8_t* __restrict__ k_q, int8_t* __restrict__ v_q,
|
||||
uint16_t* __restrict__ k_scale, uint16_t* __restrict__ v_scale,
|
||||
const int32_t* __restrict__ table, const int32_t* __restrict__ step,
|
||||
const float* __restrict__ kcur, const float* __restrict__ vcur, int kv_heads,
|
||||
int head_dim, int page_size, KvHostPools host) {
|
||||
int head_dim, int page_size, KvHostPools host, int step_stride = 0,
|
||||
int cur_stride = 0) {
|
||||
if constexpr (MULTI) {
|
||||
const int j = blockIdx.z >> 1;
|
||||
step += (size_t) j * step_stride; kcur += (size_t) j * cur_stride; vcur += (size_t) j * cur_stride;
|
||||
}
|
||||
const long long pos = (long long) __ldg(step + kStepPos);
|
||||
const int h = blockIdx.x, g = blockIdx.y, t = threadIdx.x;
|
||||
const bool is_v = blockIdx.z == 1;
|
||||
const bool is_v = MULTI ? (blockIdx.z & 1) == 1 : blockIdx.z == 1;
|
||||
const int groups = head_dim / KV_Q8_GROUP;
|
||||
const float x = (is_v ? vcur : kcur)[h * head_dim + g * KV_Q8_GROUP + t];
|
||||
// max |x| over the 64 values: two warps, then combine through shared memory in a fixed order
|
||||
@@ -111,6 +119,18 @@ void kv_append_q8_step(int8_t* k_q, int8_t* v_q, uint16_t* k_scale, uint16_t* v_
|
||||
check("kv_append_q8 launch");
|
||||
}
|
||||
|
||||
void kv_append_q8_steps(int8_t* k_q, int8_t* v_q, uint16_t* k_scale, uint16_t* v_scale, const int32_t* page_table,
|
||||
const int32_t* step, int step_stride, const float* kcur, const float* vcur, int cur_stride,
|
||||
int n_tok, const QsaShapes& s, void* stream, const KvHostPools* host) {
|
||||
validate(s, "kv_append_q8 (steps)");
|
||||
if (n_tok < 1) return;
|
||||
const dim3 grid((unsigned) s.n_head_kv, (unsigned) (s.head_dim / KV_Q8_GROUP), (unsigned) (2 * n_tok));
|
||||
kv_append_q8_kernel<true><<<grid, KV_Q8_GROUP, 0, (cudaStream_t) stream>>>(
|
||||
k_q, v_q, k_scale, v_scale, page_table, step, kcur, vcur, (int) s.n_head_kv, (int) s.head_dim,
|
||||
(int) s.page_size, host ? *host : KvHostPools{}, step_stride, cur_stride);
|
||||
check("kv_append_q8 (steps) launch");
|
||||
}
|
||||
|
||||
void kv_gather_q8_step(const int8_t* k_q, const int8_t* v_q, const uint16_t* k_scale, const uint16_t* v_scale,
|
||||
const int32_t* page_table, const int32_t* ids, const int32_t* step, int64_t max_ids,
|
||||
const QsaShapes& s, uint16_t* k_scratch, uint16_t* v_scratch, void* stream) {
|
||||
|
||||
@@ -124,11 +124,18 @@ __global__ void bf16_f32_mmvf_multi_kernel(const float* __restrict__ x, int64_t
|
||||
// rows again (n_out blocks x T x n_in floats through L2 - the part that grew with T); here a block reads the
|
||||
// activations once for RPB rows. Per output, thread t still walks pairs t, t + BLOCK_SIZE, ... with the same two
|
||||
// ordered FMAs and the same warp and block reductions: bit-identical to bf16_f32_mmvf_multi_kernel.
|
||||
template <int BLOCK_SIZE, int NT, int RPB, bool TS = false>
|
||||
// S26 STRATA_LFUSE=1 (AUX): one more block computes row 0 of w_aux into y_aux[k * ldy_aux] - the same per-output
|
||||
// code, so each aux output is bitwise what its own bf16_f32_mmvf_multi_kernel launch gave (the shared expert's gate)
|
||||
template <int BLOCK_SIZE, int NT, int RPB, bool TS = false, bool AUX = false>
|
||||
__global__ void bf16_f32_mmvf_rows_kernel(const float* __restrict__ x, int64_t ldx, const uint16_t* __restrict__ w,
|
||||
float* __restrict__ y, int64_t ldy, int n_in, int n_out, int n_tok) {
|
||||
float* __restrict__ y, int64_t ldy, int n_in, int n_out, int n_tok,
|
||||
const uint16_t* __restrict__ w_aux = nullptr, float* __restrict__ y_aux = nullptr,
|
||||
int64_t ldy_aux = 0) {
|
||||
const int t = threadIdx.x;
|
||||
const int o0 = blockIdx.x * RPB;
|
||||
int o0 = blockIdx.x * RPB;
|
||||
if constexpr (AUX) {
|
||||
if (blockIdx.x == gridDim.x - 1) { w = w_aux; y = y_aux; ldy = ldy_aux; n_out = 1; o0 = 0; }
|
||||
}
|
||||
__shared__ float partials[RPB][NT][32];
|
||||
if constexpr (BLOCK_SIZE > 32) {
|
||||
if (t < 32)
|
||||
@@ -210,10 +217,18 @@ __global__ void bf16_f32_mmvf_rows_kernel(const float* __restrict__ x, int64_t l
|
||||
|
||||
template <int B>
|
||||
void launch_rows(const float* x, int64_t ldx, const uint16_t* w, float* y, int64_t ldy, int n_in, int n_out, int n_tok,
|
||||
cudaStream_t st) {
|
||||
cudaStream_t st, const uint16_t* w_aux = nullptr, float* y_aux = nullptr, int64_t ldy_aux = 0) {
|
||||
constexpr int RPB = 4;
|
||||
const unsigned nb = (unsigned) ((n_out + RPB - 1) / RPB);
|
||||
static const bool ts = [] { const char* v = std::getenv("STRATA_TSUM"); return v && v[0] == '1'; }();
|
||||
if (w_aux != nullptr) { // S26 STRATA_LFUSE: + one aux row block
|
||||
#define S26_AUX(TSV) \
|
||||
if (n_tok <= 4) bf16_f32_mmvf_rows_kernel<B, 4, RPB, TSV, true><<<nb + 1, B, 0, st>>>(x, ldx, w, y, ldy, n_in, n_out, n_tok, w_aux, y_aux, ldy_aux); \
|
||||
else bf16_f32_mmvf_rows_kernel<B, 8, RPB, TSV, true><<<nb + 1, B, 0, st>>>(x, ldx, w, y, ldy, n_in, n_out, n_tok, w_aux, y_aux, ldy_aux);
|
||||
if (ts) { S26_AUX(true) } else { S26_AUX(false) }
|
||||
#undef S26_AUX
|
||||
return;
|
||||
}
|
||||
if (ts) {
|
||||
if (n_tok <= 4) bf16_f32_mmvf_rows_kernel<B, 4, RPB, true><<<nb, B, 0, st>>>(x, ldx, w, y, ldy, n_in, n_out, n_tok);
|
||||
else bf16_f32_mmvf_rows_kernel<B, 8, RPB, true><<<nb, B, 0, st>>>(x, ldx, w, y, ldy, n_in, n_out, n_tok);
|
||||
@@ -238,6 +253,32 @@ int mmvf_block_size(int64_t n_in) {
|
||||
|
||||
} // namespace
|
||||
|
||||
bool bf16_gemv_fp32_mmvf_multi_aux(const float* x, int64_t ldx, const uint16_t* w, float* y, int64_t ldy,
|
||||
int64_t n_in, int64_t n_out, int n_tok, const uint16_t* w_aux, float* y_aux,
|
||||
int64_t ldy_aux, void* stream) {
|
||||
static const bool rows = [] { const char* v = std::getenv("STRATA_MMVF_ROWS"); return v && v[0] == '1'; }();
|
||||
// exactly the conditions under which both calls would take the rows / multi kernels at the same block size
|
||||
if (!rows || n_out < 64 || n_tok < 2 || n_tok > 8 || n_in <= 0 || (n_in & 1) != 0 || (ldx & 1) != 0 || x == nullptr ||
|
||||
w == nullptr || y == nullptr || w_aux == nullptr || y_aux == nullptr || (reinterpret_cast<uintptr_t>(x) & 7u) != 0)
|
||||
return false;
|
||||
const cudaStream_t st = (cudaStream_t) stream;
|
||||
const int ni = (int) n_in, no = (int) n_out;
|
||||
switch (mmvf_block_size(n_in)) {
|
||||
case 32: launch_rows<32>(x, ldx, w, y, ldy, ni, no, n_tok, st, w_aux, y_aux, ldy_aux); break;
|
||||
case 64: launch_rows<64>(x, ldx, w, y, ldy, ni, no, n_tok, st, w_aux, y_aux, ldy_aux); break;
|
||||
case 96: launch_rows<96>(x, ldx, w, y, ldy, ni, no, n_tok, st, w_aux, y_aux, ldy_aux); break;
|
||||
case 128: launch_rows<128>(x, ldx, w, y, ldy, ni, no, n_tok, st, w_aux, y_aux, ldy_aux); break;
|
||||
case 160: launch_rows<160>(x, ldx, w, y, ldy, ni, no, n_tok, st, w_aux, y_aux, ldy_aux); break;
|
||||
case 192: launch_rows<192>(x, ldx, w, y, ldy, ni, no, n_tok, st, w_aux, y_aux, ldy_aux); break;
|
||||
case 224: launch_rows<224>(x, ldx, w, y, ldy, ni, no, n_tok, st, w_aux, y_aux, ldy_aux); break;
|
||||
default: launch_rows<256>(x, ldx, w, y, ldy, ni, no, n_tok, st, w_aux, y_aux, ldy_aux); break;
|
||||
}
|
||||
const cudaError_t result = cudaGetLastError();
|
||||
if (result != cudaSuccess)
|
||||
throw std::runtime_error(std::string("bf16_gemv_fp32_mmvf_multi_aux launch: ") + cudaGetErrorString(result));
|
||||
return true;
|
||||
}
|
||||
|
||||
void bf16_gemv_fp32_mmvf_multi(const float* x, int64_t ldx, const uint16_t* w, float* y, int64_t ldy,
|
||||
int64_t n_in, int64_t n_out, int n_tok, void* stream) {
|
||||
if (n_tok == 1 && ldy >= n_out) { bf16_gemv_fp32_mmvf(x, w, y, n_in, n_out, stream); return; }
|
||||
|
||||
@@ -1009,14 +1009,21 @@ static bool s26_tsum_on() {
|
||||
return on;
|
||||
}
|
||||
|
||||
template<typename F, int NCOLS, int NW, int ROWS, bool TS = false>
|
||||
// S26 STRATA_LFUSE (PAIR): blocks n_out.. compute the same rows of w2 into y2 (two matrices of one shape on one
|
||||
// input in one launch; each output's code is the single matrix's)
|
||||
template<typename F, int NCOLS, int NW, int ROWS, bool TS = false, bool PAIR = false>
|
||||
__launch_bounds__(NW * WARP, 1)
|
||||
__global__ void native_mmvq_multi_kernel(const typename F::Block* __restrict__ w,
|
||||
const Q81Block* __restrict__ x,
|
||||
float* __restrict__ y, int n_in, int n_out) {
|
||||
float* __restrict__ y, int n_in, int n_out,
|
||||
const typename F::Block* __restrict__ w2 = nullptr, float* __restrict__ y2 = nullptr) {
|
||||
constexpr int BPI = F::BPI * NW / WARPS; // blocks per iteration scale with the warp count
|
||||
const int tid = WARP * int(threadIdx.y) + int(threadIdx.x);
|
||||
const int row0 = ROWS * int(blockIdx.x);
|
||||
int bxi = int(blockIdx.x);
|
||||
if constexpr (PAIR) {
|
||||
if (bxi >= n_out / ROWS) { bxi -= n_out / ROWS; w = w2; y = y2; }
|
||||
}
|
||||
const int row0 = ROWS * bxi;
|
||||
const int blocks_per_row = n_in / F::DIV;
|
||||
const int x_stride = n_in / Q8K; // Q8_1 blocks per activation column
|
||||
float tmp[NCOLS][ROWS] = {};
|
||||
@@ -1874,6 +1881,43 @@ void native_q5_0_f32(const void* weights, const float* x, void* scratch_q8_1,
|
||||
small_f32<Q50Block, 4>(weights, x, scratch_q8_1, y, n_in, n_out, ncols, stream);
|
||||
}
|
||||
|
||||
template<int NCOLS>
|
||||
void q8_0_pair_launch(const Q80Block* w1, const Q80Block* w2, const Q81Block* x, float* y1, float* y2, int n_in, int n_out,
|
||||
cudaStream_t s) {
|
||||
using F = SmallTraits<Q80Block, 8>;
|
||||
const dim3 threads(WARP, WARPS);
|
||||
if (s26_tsum_on())
|
||||
native_mmvq_multi_kernel<F, NCOLS, WARPS, 1, true, true><<<unsigned(2 * n_out), threads, 0, s>>>(w1, x, y1, n_in, n_out, w2, y2);
|
||||
else
|
||||
native_mmvq_multi_kernel<F, NCOLS, WARPS, 1, false, true><<<unsigned(2 * n_out), threads, 0, s>>>(w1, x, y1, n_in, n_out, w2, y2);
|
||||
}
|
||||
|
||||
bool native_mmvq_pair(int ggml_type, const void* w1, const void* w2, const void* x_q8_1, float* y1, float* y2,
|
||||
int n_in, int n_out, int ncols, void* stream) {
|
||||
using F = SmallTraits<Q80Block, 8>;
|
||||
// only where both calls would run native_mmvq_multi_kernel<F, ncols, WARPS, 1> (exact layout, not packed)
|
||||
if (ggml_type != 8 || !g_multi_exact || ncols < 2 || ncols > 8 || n_out <= 0 || n_in <= 0 || n_in % 32 != 0 ||
|
||||
n_in / F::DIV < F::BPI || !w1 || !w2 || !x_q8_1 || !y1 || !y2 || !stream)
|
||||
return false;
|
||||
const auto& reg = q8_packed_registry();
|
||||
if (reg.count(w1) || reg.count(w2)) return false;
|
||||
const auto* a = static_cast<const Q80Block*>(w1);
|
||||
const auto* b = static_cast<const Q80Block*>(w2);
|
||||
const auto* x = static_cast<const Q81Block*>(x_q8_1);
|
||||
const auto s = static_cast<cudaStream_t>(stream);
|
||||
switch (ncols) {
|
||||
case 2: q8_0_pair_launch<2>(a, b, x, y1, y2, n_in, n_out, s); break;
|
||||
case 3: q8_0_pair_launch<3>(a, b, x, y1, y2, n_in, n_out, s); break;
|
||||
case 4: q8_0_pair_launch<4>(a, b, x, y1, y2, n_in, n_out, s); break;
|
||||
case 5: q8_0_pair_launch<5>(a, b, x, y1, y2, n_in, n_out, s); break;
|
||||
case 6: q8_0_pair_launch<6>(a, b, x, y1, y2, n_in, n_out, s); break;
|
||||
case 7: q8_0_pair_launch<7>(a, b, x, y1, y2, n_in, n_out, s); break;
|
||||
default: q8_0_pair_launch<8>(a, b, x, y1, y2, n_in, n_out, s); break;
|
||||
}
|
||||
launch_check();
|
||||
return true;
|
||||
}
|
||||
|
||||
void native_q8_0_mmvq(const void* weights, const void* x_q8_1, float* y,
|
||||
int n_in, int n_out, int ncols, void* stream) {
|
||||
if (q8_packed_mmvq(weights, x_q8_1, y, n_in, n_out, ncols, stream)) return; // STRATA_Q8_PACKED=1 tensors
|
||||
|
||||
@@ -31,10 +31,13 @@
|
||||
namespace strata::kernels {
|
||||
namespace {
|
||||
std::atomic<bool> enabled{false};
|
||||
template<bool ZADD = false>
|
||||
// GATED (S26 STRATA_LFUSE): `shared` is the shared expert's unscaled output and sg its raw gate logit per token;
|
||||
// the row is scaled here as shared_expert_multi did (native_scalar_sigmoid_multi_kernel's expression, then one
|
||||
// rounded multiply = scale_rows_kernel's `out *= g`), then added as before
|
||||
template<bool ZADD = false, bool GATED = false>
|
||||
__global__ void combine(const float* __restrict__ parts, const float* __restrict__ weights,
|
||||
const float* __restrict__ shared, float* __restrict__ output,
|
||||
int64_t n_embd, int k) {
|
||||
int64_t n_embd, int k, const float* __restrict__ sg = nullptr) {
|
||||
// blockIdx.y = the token of a multi-token launch (0 for the single one)
|
||||
const int64_t tk = blockIdx.y;
|
||||
parts += tk * k * n_embd; weights += tk * k; if (shared) shared += tk * n_embd; output += tk * n_embd;
|
||||
@@ -49,7 +52,15 @@ __global__ void combine(const float* __restrict__ parts, const float* __restrict
|
||||
if (ZADD) p = 0.0f + p;
|
||||
sum += p * weights[expert];
|
||||
}
|
||||
if (shared) sum += shared[col];
|
||||
if constexpr (GATED) {
|
||||
// contraction off in this block only: the product is rounded on its own (scale_rows' store), then added
|
||||
#pragma clang fp contract(off)
|
||||
const float g = __fdividef(1.0f, 1.0f + __expf(-sg[tk]));
|
||||
const float sh = shared[col] * g;
|
||||
sum += sh;
|
||||
} else {
|
||||
if (shared) sum += shared[col];
|
||||
}
|
||||
output[col] = sum;
|
||||
}
|
||||
bool valid_span(const void* p, size_t bytes) {
|
||||
@@ -89,6 +100,16 @@ void native_moe_combine_multi(const float* parts, const float* weights, const fl
|
||||
const auto error = cudaGetLastError();
|
||||
if (error != cudaSuccess) throw std::runtime_error(cudaGetErrorString(error));
|
||||
}
|
||||
void native_moe_combine_multi_hits_gated(const float* hits, const float* weights, const float* shared,
|
||||
const float* shared_gate, float* output, int64_t n_embd, int64_t k, int n_tok,
|
||||
void* stream) {
|
||||
if (!stream || n_embd <= 0 || k < 1 || k > 15 || n_tok < 1 || !shared || !shared_gate)
|
||||
throw std::invalid_argument("native MoE combine (hits, gated) requires a stream, width, 1..15 experts, tokens");
|
||||
combine<true, true><<<dim3(unsigned((n_embd + 255) / 256), (unsigned) n_tok), 256, 0, static_cast<cudaStream_t>(stream)>>>(
|
||||
hits, weights, shared, output, n_embd, int(k), shared_gate);
|
||||
const auto error = cudaGetLastError();
|
||||
if (error != cudaSuccess) throw std::runtime_error(cudaGetErrorString(error));
|
||||
}
|
||||
void native_moe_combine_multi_hits(const float* hits, const float* weights, const float* shared, float* output,
|
||||
int64_t n_embd, int64_t k, int n_tok, void* stream) {
|
||||
if (!stream || n_embd <= 0 || k < 1 || k > 15 || n_tok < 1)
|
||||
|
||||
@@ -191,14 +191,19 @@ __global__ void scale_rows_kernel(float* __restrict__ out, const float* __restri
|
||||
|
||||
void shared_expert_multi(int n_tok, const float* x, const uint16_t* x_bf16, const NativeSharedWeights& nw,
|
||||
const uint16_t* gate_inp_bf16, float* gate, float* up, float* g, float* out, int64_t n_embd,
|
||||
int64_t n_ff, void* stream, const void* x_q8_1_ready) {
|
||||
int64_t n_ff, void* stream, const void* x_q8_1_ready, int lfuse) {
|
||||
const bool gate_deferred = (lfuse & 1) != 0, pair = (lfuse & 2) != 0;
|
||||
if (n_tok < 1 || n_tok > 8 || !nw.q8_1 || !nw.gate_data || !nw.up_data || !nw.down_data || !stream)
|
||||
throw std::invalid_argument("shared_expert_multi: needs 1..8 tokens, native weights, scratch and a stream");
|
||||
cudaStream_t cs = (cudaStream_t) stream;
|
||||
const void* xq = x_q8_1_ready ? x_q8_1_ready : nw.q8_1;
|
||||
if (!x_q8_1_ready) native_quantize_q8_1(x, nw.q8_1, (int) n_embd, n_tok, stream);
|
||||
native_mmvq(nw.gate_type, nw.gate_data, xq, gate, (int) n_embd, (int) n_ff, n_tok, stream);
|
||||
native_mmvq(nw.up_type, nw.up_data, xq, up, (int) n_embd, (int) n_ff, n_tok, stream);
|
||||
// S26 STRATA_LFUSE (gate_deferred): gate and up in one launch where the pair applies (bitwise the two calls)
|
||||
if (!(pair && nw.gate_type == nw.up_type &&
|
||||
native_mmvq_pair(nw.gate_type, nw.gate_data, nw.up_data, xq, gate, up, (int) n_embd, (int) n_ff, n_tok, stream))) {
|
||||
native_mmvq(nw.gate_type, nw.gate_data, xq, gate, (int) n_embd, (int) n_ff, n_tok, stream);
|
||||
native_mmvq(nw.up_type, nw.up_data, xq, up, (int) n_embd, (int) n_ff, n_tok, stream);
|
||||
}
|
||||
const int n = (int) (n_ff * n_tok);
|
||||
if (x_q8_1_ready != nullptr && (n_ff % 32) == 0) { // STRATA_VERIFY_QDEDUP: swiglu + its q8_1 in one launch
|
||||
native_swiglu_q8_1_kernel<<<(unsigned) ((n + 255) / 256), 256, 0, cs>>>(gate, up, gate, (SwqQ81*) nw.q8_1, n);
|
||||
@@ -207,6 +212,11 @@ void shared_expert_multi(int n_tok, const float* x, const uint16_t* x_bf16, cons
|
||||
native_quantize_q8_1(gate, nw.q8_1, (int) n_ff, n_tok, stream);
|
||||
}
|
||||
native_mmvq(nw.down_type, nw.down_data, nw.q8_1, out, (int) n_ff, (int) n_embd, n_tok, stream);
|
||||
if (gate_deferred) { // the caller computed g (raw) with the router and scales `out` in the combine
|
||||
const cudaError_t e = cudaGetLastError();
|
||||
if (e != cudaSuccess) throw std::runtime_error(std::string("shared_expert_multi: ") + cudaGetErrorString(e));
|
||||
return;
|
||||
}
|
||||
static const bool batch = [] { const char* v = std::getenv("STRATA_DEC_BATCH"); return v == nullptr || std::atoi(v) != 0; }();
|
||||
if (native_bf16 && batch && n_tok > 1) { // one gemv for all rows (outputs identical), one sigmoid launch
|
||||
bf16_gemv_fp32_mmvf_multi(x, n_embd, gate_inp_bf16, g, 1, n_embd, 1, n_tok, stream);
|
||||
|
||||
Reference in new issue
Block a user