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

IQ4_XS 5-column interleaved MMVQ runs as 6 columns with 5 stored (STRATA_IL_PAD, bitwise); --draft auto: MTP requests keep the MTP's prompt-lookup width

Browse Source
This commit is contained in:
Niko1221 committed 2026-11-01 03:28:12 +01:00
1 parent 1db96cda5f
commit 4c5a554912
2 files changed
+26 -8
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

src/kernels/cuda/native_mmvq.cu
+22 -6
View File
@@ -2151,7 +2151,9 @@ __global__ void native_mmvq_il_pf_kernel(const typename I::F::Block* __restrict_
}
}
template <typename I, int NC, int R>
// NOUT < NC (BW-NEXT, STRATA_IL_PAD): the kernel of NC columns computes a column that is not wanted and stores only the first NOUT; each column's value is
// its own sum, so the stored ones are bitwise those of the NOUT-column kernel. 5 columns of IQ4_XS run as 6 (the 5-column instantiation is 15-40% slower).
template <typename I, int NC, int R, int NOUT = NC>
__launch_bounds__(4 * WARP)
__global__ void native_mmvq_il_kernel(const typename I::F::Block* __restrict__ w, const int* __restrict__ xq,
const float* __restrict__ xd, float* __restrict__ y, int n_in, int n_out) {
@@ -2199,7 +2201,7 @@ __global__ void native_mmvq_il_kernel(const typename I::F::Block* __restrict__ w
s += acc[2][i][c];
s += acc[3][i][c];
s = warp_sum(s);
if (lane == 0 && row0 + i < n_out) y[std::size_t(c) * n_out + row0 + i] = s;
if (c < NOUT && lane == 0 && row0 + i < n_out) y[std::size_t(c) * n_out + row0 + i] = s;
}
}
}
@@ -2222,10 +2224,10 @@ int il_pf_for(int type, int ncols, int n_out, int* rows) {
return t >> 4;
}
// one (type, columns, rows) launch: today's kernel, or the variant il_pf_for names for a type that has them
template <typename I, int NC, int R>
template <typename I, int NC, int R, int NOUT = NC>
void launch_il_one(unsigned blocks, const typename I::F::Block* w, const int* xq, const float* xd, float* y, int n_in, int n_out,
cudaStream_t s, int pf) {
if constexpr (I::HAS_PF) {
if constexpr (I::HAS_PF && NOUT == NC) {
switch (pf) {
case 6: native_mmvq_il_pf_kernel<I, NC, R, 0, 4, false><<<blocks, 4 * WARP, 0, s>>>(w, xq, xd, y, n_in, n_out); return;
case 2: native_mmvq_il_pf_kernel<I, NC, R, 2, 4, true><<<blocks, 4 * WARP, 0, s>>>(w, xq, xd, y, n_in, n_out); return;
@@ -2235,7 +2237,13 @@ void launch_il_one(unsigned blocks, const typename I::F::Block* w, const int* xq
default: break;
}
}
native_mmvq_il_kernel<I, NC, R><<<blocks, 4 * WARP, 0, s>>>(w, xq, xd, y, n_in, n_out);
native_mmvq_il_kernel<I, NC, R, NOUT><<<blocks, 4 * WARP, 0, s>>>(w, xq, xd, y, n_in, n_out);
}
// STRATA_IL_PAD (bit 0: IQ4_XS 5 columns as 6; default on, 0 off): bitwise, see native_mmvq_il_kernel
int il_pad() {
static const int v = [] { const char* e = std::getenv("STRATA_IL_PAD"); return e != nullptr ? std::atoi(e) : 1; }();
return v;
}
template <typename I, int R>
@@ -2255,8 +2263,16 @@ void launch_il(const void* weights, const void* x_il, float* y, int n_in, int n_
#undef STRATA_IL_CASE
// 147-verify: 5..8 columns (the cp = 8 copy) at 1 or 2 rows a warp; launch_il_rows maps 4 rows to 2 there
#define STRATA_IL_CASE(N) case N: if constexpr (R <= 2) launch_il_one<I, N, R>(blocks, w, xq, xd, y, n_in, n_out, s, pf); break;
STRATA_IL_CASE(5) STRATA_IL_CASE(6) STRATA_IL_CASE(7) STRATA_IL_CASE(8)
STRATA_IL_CASE(6) STRATA_IL_CASE(7) STRATA_IL_CASE(8)
#undef STRATA_IL_CASE
case 5:
if constexpr (R <= 2) {
if constexpr (std::is_same_v<I, IlIQ4XS>) {
if ((il_pad() & 1) != 0) { launch_il_one<I, 6, R, 5>(blocks, w, xq, xd, y, n_in, n_out, s, 0); break; }
}
launch_il_one<I, 5, R>(blocks, w, xq, xd, y, n_in, n_out, s, pf);
}
break;
default: throw std::invalid_argument("native_mmvq_il requires 1 <= ncols <= 8");
}
}
src/program/generate.cpp
+4 -2
View File
@@ -12434,6 +12434,7 @@ int main(int argc, char** argv) {
// --draft auto: a request that drafts with the MTP layer has its own widest window (the MTP loses to itself above 4 on Blackwell), the DFlash2 requests S_mtp_wide
const int S_mtp_narrow = o.draft_auto ? std::min(S_mtp_wide, std::max(2, o.draft_mtp_spec)) : S_mtp_wide;
int S_mtp = S_mtp_wide; // this request's (set when it starts)
int S_look = S; // --draft auto: the widest window a prompt-lookup draft may make for this request (the MTP requests keep the MTP's width, as without DFlash2)
if (use_mtp && S_mtp < S) mtp.set_max_drafts(S_mtp - 1);
// --pipeline-windows 2 loads the drafter with 8 rows: the serial loop's chain keeps its own length
else if (use_mtp && mtp.max_t() > S_mtp) mtp.set_max_drafts(S_mtp - 1);
@@ -14096,6 +14097,7 @@ int main(int argc, char** argv) {
if (!req_df2 && !req_min_p_set && o.draft_mtp_min_p >= 0.0) req_spec_min_p = o.draft_mtp_min_p;
if (req_df2 != policy_is_df2) { std::swap(policy, policy_idle); policy_is_df2 = req_df2; }
S_mtp = req_df2 ? S_mtp_wide : S_mtp_narrow;
S_look = req_df2 ? S : std::min(S, S_mtp_narrow);
if (use_mtp) mtp.set_max_drafts(S_mtp - 1);
}
if (req_prefill_only && (geni || o.batch > 0 || n < 2)) {
@@ -16941,7 +16943,7 @@ int main(int argc, char** argv) {
// a repeat of earlier context (prompt lookup) where the MTP's own first guess agrees: the policy takes it
// when its expected tokens per ms, from the measured acceptance and window costs, beat the MTP window's
if (o.suffix_draft > 0 && !first_window) {
const int k = sfx.propose(S - 1, sbuf.data());
const int k = sfx.propose(S_look - 1, sbuf.data());
sfx_match = sfx.last_match();
// a probabilistic window must not condition the source on the draw: the policy copy fixed at the request's start
// chooses, and the lookup is taken without asking whether the MTP's first guess agrees
@@ -16960,7 +16962,7 @@ int main(int argc, char** argv) {
int csrc = -1;
const int nt = chain_tail(sfx, drafts.data(), T - 1, ctail);
chain_n = strata::spec::propose_from_sources(lookup_src, ctail.data(), nt, T - 1,
std::min(o.lookup_chain, S - T), o.lookup_chain_min,
std::min(o.lookup_chain, S_look - T), o.lookup_chain_min,
cbuf.data(), &cm, &csrc);
// the policy keeps the chained tokens that pay at this machine's measured window costs
if (chain_n > 0 && !chain_fixed) {