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:
Niko1221 committed 2026-10-04 00:41:01 +00:00
1 parent 75c9fc6baa
commit b08cf3e160
11 files changed
+204 -17

No files matched your search

+6
View File
@@ -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
+5
View File
@@ -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);
+4
View File
@@ -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
+5
View File
@@ -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);
}
+2 -1
View File
@@ -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
View File
@@ -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 {
+22 -2
View File
@@ -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) {
+45 -4
View File
@@ -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; }
+47 -3
View File
@@ -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
+24 -3
View File
@@ -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)
+13 -3
View File
@@ -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);