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

Remove STRATA_DF2_COUPLED (coupled sampled picks): at a fixed 8-row window it does not raise acceptance (tokens per round -1.7% short / -0.4% long, first-draft match 0.616 vs 0.629; draft temperature 0.6 / 1.5 / 2.5: 2.885 / 2.728 / 2.559 vs 2.868 tokens per round plain), the selector's argmax is the better draft

Browse Source
This commit is contained in:
Niko1221 committed 2026-11-01 15:02:51 +01:00
1 parent 5f24b96dbd
commit a84eebb513
5 files changed
+19 -138
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

include/strata/core/dflash2.hpp
-13
View File
@@ -17,7 +17,6 @@
#pragma once
#include "strata/core/session.hpp"
#include "strata/kernels/sampler.hpp"
#include <cuda_runtime.h>
@@ -35,7 +34,6 @@ class Dflash2Drafter {
public:
static constexpr int kRows = 8; ///< block rows: the anchor + 7 masks
static constexpr int kDrafts = kRows - 1;
static constexpr int kMeta = 64; ///< int32 of the staging block: [0..31] positions and tokens, [32..39] the coupled picks (kernels/dflash2.hpp df2_select)
static constexpr int kRing = 2064; ///< K/V slots per layer: the 2048 window + slack for rows written ahead of a rejection
Dflash2Drafter() = default;
@@ -70,13 +68,6 @@ public:
/// selector's pick and its probability). The taps pointer must be the same device buffer on every call (graphs are captured on it).
bool draft(const float* taps_dev, int n_new, int64_t pos0, int32_t anchor, int32_t* drafts, float* probs, std::string& err);
/// COUPLED PICKS (STRATA_DF2_COUPLED=1, off by default; core/coupled_draft.hpp is the MTP's version). For a sampled request the selector's walk picks each draft as the
/// target's sampler would from the selector's 16 scores, with the uniform the verify draws for the row that checks it, instead of the argmax. Verification stays
/// exact-match, so the text is the plain seeded sampler's whatever the drafts are; the drafts only agree more often. Off for greedy requests, penalties and logit
/// bias (the picks are the argmax then). STRATA_DF2_COUPLED_DT=x runs the draft's chain at x times the request's temperature.
void set_sampling(const strata::kernels::SamplerParams& sp);
bool coupled() const { return coupled_on_; }
/// Captures the eight round graphs (n_new = 1..8) on `taps` now, so no round pays for a capture.
bool prepare(const float* taps_dev, std::string& err);
@@ -170,10 +161,6 @@ private:
int64_t* tag_ = nullptr; ///< per ring slot: (epoch << 32 | position) of the row it holds, -1 none
cudaEvent_t ev_a_ = nullptr, ev_b_ = nullptr;
bool timing_ = false;
bool coupled_on_ = false;
uint64_t c_seed_ = 0;
int32_t c_topk_ = 0;
float c_temp_ = 1.0f, c_topp_ = 1.0f, c_minp_ = 0.0f;
};
} // namespace strata::core
include/strata/kernels/dflash2.hpp
+1 -5
View File
@@ -91,12 +91,8 @@ void df2_topk(const float* logits, int n, const int32_t* idx_map, int rows, uint
/// The candidate selector's walk: for row i = 0..rows-1 the 16 candidates' scores are unary[i][j] + dot(succ[tok[i][j]],
/// pred[prev] * hp[i]) over `rank` values (codebooks are bf16 rows of `rank`), the pick is the argmax, prev = the pick. Writes
/// out_tok[i] and out_prob[i] (the softmax probability of the pick among the 16). `prev0` is read from *anchor_tok on the device.
/// `samp` (device, or null): COUPLED PICKS (STRATA_DF2_COUPLED) - 8 int32: [0] 1 = on, [1] [2] the request's Philox seed (low, high), [3] top_k, [4] the draft
/// temperature (float bits), [5] top_p, [6] min_p (float bits), [7] the anchor's position. Row i then picks as the target's sampler would from these 16 scores
/// (top_k, top_p on the raw scores, min_p, temperature, inverse CDF) with the uniform Philox(seed, anchor position + i) that the verify row checking draft i draws,
/// and out_prob[i] is the pick's probability in that distribution.
void df2_select(const float* hp, const float* unary, const int32_t* tok, const uint16_t* pred_cb, const uint16_t* succ_cb, int rank, int rows,
const int32_t* anchor_tok, int32_t* out_tok, float* out_prob, void* stream, const int32_t* samp = nullptr);
const int32_t* anchor_tok, int32_t* out_tok, float* out_prob, void* stream);
/// The verify window's tap: out[r][d] = mean over the hc streams c of (R[r][c][d] + bo[r][d] * 2 * sigmoid(inject[r][c] / hc)) - the layer output with
/// its pending FFN write (gr_write) applied, averaged over the streams. R rows are `R_stride` floats apart, bo rows `bo_stride`, inject rows
src/core/dflash2.cpp
+7 -31
View File
@@ -238,7 +238,7 @@ bool Dflash2Drafter::alloc(std::string& err) {
return true;
};
const size_t B = kRows;
bool ok = dev(&d_meta_, kMeta) && dev(&xq_, kernels::native_q8_1_bytes(N_, kRows) + 4096) && dev(&xil_, kernels::native_q8_1_il_bytes(N_, kRows) + 4096) && dev(&xqp_, (size_t) 5 * N_ * 8 + 4096) && dev(&xdp_, (size_t) 5 * N_ / 32 * 8 + 64) && dev(&xqc_, (size_t) 5 * N_ * 8 + 4096) && dev(&xdc_, (size_t) 5 * N_ / 32 * 8 + 64) && dev(&rt_ctx_, 2 * B * 128) && dev(&rope_inv_, 64) &&
bool ok = dev(&d_meta_, 32) && dev(&xq_, kernels::native_q8_1_bytes(N_, kRows) + 4096) && dev(&xil_, kernels::native_q8_1_il_bytes(N_, kRows) + 4096) && dev(&xqp_, (size_t) 5 * N_ * 8 + 4096) && dev(&xdp_, (size_t) 5 * N_ / 32 * 8 + 64) && dev(&xqc_, (size_t) 5 * N_ * 8 + 4096) && dev(&xdc_, (size_t) 5 * N_ / 32 * 8 + 64) && dev(&rt_ctx_, 2 * B * 128) && dev(&rope_inv_, 64) &&
dev(&fcy_, B * N_) && dev(&hctx_, B * N_) && dev(&kc_, B * NKV_ * HD_) && dev(&vc_, B * NKV_ * HD_) && dev(&h_, B * N_) && dev(&xn_, B * N_) &&
dev(&xc_, B * N_) && dev(&dyn_, B * 640) && dev(&q_, B * NH_ * HD_) && dev(&qn_, B * NH_ * HD_) && dev(&k_, B * NKV_ * HD_) && dev(&v_, B * NKV_ * HD_) &&
dev(&kblk_, B * NKV_ * HD_) && dev(&ao_, B * NH_ * HD_) && dev(&o_, B * N_) && dev(&gate_, B * I_) && dev(&up_, B * I_) && dev(&mlp_, B * N_) &&
@@ -248,10 +248,9 @@ bool Dflash2Drafter::alloc(std::string& err) {
if (!ok) { err = "dflash2: the drafter's buffers do not fit in VRAM"; return false; }
void* hd = nullptr;
void* dd = nullptr;
if (!mapped(kMeta * 4, (void**) &hd, &dd)) { err = "dflash2: mapped staging"; return false; }
if (!mapped(32 * 4, (void**) &hd, &dd)) { err = "dflash2: mapped staging"; return false; }
h_meta_ = (int32_t*) hd;
m_meta_ = (int32_t*) dd;
std::memset(h_meta_, 0, kMeta * 4);
if (!mapped(64 * 4, &hd, &dd)) { err = "dflash2: mapped output"; return false; }
h_out_ = (int32_t*) hd;
out_tok_ = (int32_t*) dd; // the select kernel writes the 7 picks here (host-visible) ...
@@ -458,7 +457,7 @@ bool Dflash2Drafter::record_block(cudaStream_t cs, std::string& err, bool rope)
stamp("gemm:candidate_selector.hidden_projection");
df2_select(hsel_, topval_, toptok_, reinterpret_cast<const uint16_t*>(blob_ + find("candidate_selector.predecessor_codebook").off),
reinterpret_cast<const uint16_t*>(blob_ + find("candidate_selector.successor_codebook").off), rank_, ndraft_, d_meta_ + 16, out_tok_,
out_prob_, cs, d_meta_ + 32);
out_prob_, cs);
} catch (const std::exception& e) {
err = e.what();
return false;
@@ -469,7 +468,7 @@ bool Dflash2Drafter::record_block(cudaStream_t cs, std::string& err, bool rope)
bool Dflash2Drafter::capture(int n, const float* taps, std::string& err) {
if (exec_[n]) return true;
if (cudaStreamBeginCapture(cs_, cudaStreamCaptureModeThreadLocal) != cudaSuccess) { err = "dflash2: begin capture"; return false; }
kernels::copy_i32_from_mapped(d_meta_, m_meta_, kMeta, cs_);
kernels::copy_i32_from_mapped(d_meta_, m_meta_, 32, cs_);
kernels::df2_rope_table_inv(d_meta_, 16, rope_inv_, rt_ctx_, cs_); // the context rows' and the block rows' tables in one launch, before the branches split
// the context update (fc, the keys and values) and the block's front run side by side: they meet before each layer's attention
// measured: -0.3 ms a round on an RTX 3060, +0.0..0.7 on a Strix Halo (gfx1151), so the HIP builds keep one stream (STRATA_DF2_FORK=1 / 0 overrides)
@@ -560,7 +559,7 @@ bool Dflash2Drafter::context(const float* taps, int64_t n, int64_t pos0, std::st
for (int i = 0; i < m; ++i) h_meta_[i] = (int32_t) (pos0 + r0 + i);
h_meta_[18] = epoch_;
std::atomic_thread_fence(std::memory_order_seq_cst);
kernels::copy_i32_from_mapped(d_meta_, m_meta_, kMeta, cs_);
kernels::copy_i32_from_mapped(d_meta_, m_meta_, 32, cs_);
if (!record_context(taps + (size_t) r0 * (size_t) tap_n, m, cs_, err)) return false;
if (cudaStreamSynchronize(cs_) != cudaSuccess) { err = std::string("dflash2: ") + cudaGetErrorString(cudaGetLastError()); return false; }
}
@@ -576,7 +575,7 @@ bool Dflash2Drafter::context_async(const float* taps, int64_t n, int64_t pos0, s
for (int i = 0; i < (int) n; ++i) h_meta_[i] = (int32_t) (pos0 + i);
h_meta_[18] = epoch_;
std::atomic_thread_fence(std::memory_order_seq_cst);
kernels::copy_i32_from_mapped(d_meta_, m_meta_, kMeta, cs_);
kernels::copy_i32_from_mapped(d_meta_, m_meta_, 32, cs_);
if (!record_context(taps, (int) n, cs_, err)) return false;
ctx_pending_ = true;
have_ctx_ = std::max<int64_t>(have_ctx_, pos0 + n);
@@ -612,16 +611,6 @@ bool Dflash2Drafter::draft(const float* taps, int n_new, int64_t pos0, int32_t a
h_meta_[17] = (int32_t) s;
h_meta_[24] = anchor;
for (int i = 1; i < 8; ++i) h_meta_[24 + i] = mask_token_;
h_meta_[32] = coupled_on_ ? 1 : 0;
if (coupled_on_) {
h_meta_[33] = (int32_t) (uint32_t) c_seed_;
h_meta_[34] = (int32_t) (uint32_t) (c_seed_ >> 32);
h_meta_[35] = c_topk_;
std::memcpy(&h_meta_[36], &c_temp_, 4);
std::memcpy(&h_meta_[37], &c_topp_, 4);
std::memcpy(&h_meta_[38], &c_minp_, 4);
h_meta_[39] = (int32_t) s; // the anchor's position: draft i is checked by the verify row that draws Philox(seed, s + i)
}
for (int i = 0; i < kDrafts; ++i) { ((volatile int32_t*) h_out_)[i] = -1; }
std::atomic_thread_fence(std::memory_order_seq_cst);
if (timing_) cudaEventRecord(ev_a_, cs_);
@@ -630,7 +619,7 @@ bool Dflash2Drafter::draft(const float* taps, int n_new, int64_t pos0, int32_t a
for (cudaEvent_t e : prof_->ev) cudaEventDestroy(e); // the stamps of context() calls are not a round
prof_->ev.clear();
prof_->name.clear();
kernels::copy_i32_from_mapped(d_meta_, m_meta_, kMeta, cs_);
kernels::copy_i32_from_mapped(d_meta_, m_meta_, 32, cs_);
stamp("start");
if (!record_context(taps, n_new, cs_, err) || !record_block(cs_, err)) return false;
prof_collect();
@@ -655,19 +644,6 @@ bool Dflash2Drafter::draft(const float* taps, int n_new, int64_t pos0, int32_t a
return true;
}
void Dflash2Drafter::set_sampling(const strata::kernels::SamplerParams& sp) {
static const bool env = [] { const char* e = std::getenv("STRATA_DF2_COUPLED"); return e != nullptr && *e != '\0' && std::strcmp(e, "0") != 0; }();
static const float dt = [] { const char* e = std::getenv("STRATA_DF2_COUPLED_DT"); const float v = e != nullptr ? (float) std::atof(e) : 1.0f; return v > 0.0f ? v : 1.0f; }();
const bool pen = sp.penalty_last_n > 0 && (sp.penalty_repeat != 1.0f || sp.penalty_freq != 0.0f || sp.penalty_present != 0.0f);
coupled_on_ = env && !sp.greedy && sp.temperature > 0.0f && !pen && sp.logit_bias == nullptr && !sp.gumbel;
if (!coupled_on_) return;
c_seed_ = sp.seed;
c_topk_ = sp.top_k;
c_temp_ = sp.temperature * dt;
c_topp_ = sp.top_p;
c_minp_ = sp.min_p;
}
bool Dflash2Drafter::debug_rows(std::vector<float>& rows, std::string& err) {
rows.resize((size_t) kRows * (size_t) N_);
if (cudaMemcpy(rows.data(), fn_, rows.size() * 4, cudaMemcpyDeviceToHost) != cudaSuccess) { err = "dflash2: debug copy"; return false; }
src/kernels/cuda/dflash2.cu
+11 -83
View File
@@ -777,32 +777,9 @@ __global__ void topk_merge_kernel(const float* __restrict__ cval, const int32_t*
}
// ---- the selector's walk: one block of 256 threads
// Philox 4x32-10 and its uniform: a copy of sampler.cu's (the verify's sampler draws with the same function; the coupled pick below must give the
// same float for the same (seed, counter)). dflash2_coupled_parity pins the two together.
__device__ __forceinline__ float df2_philox_uniform(uint64_t seed, uint64_t counter) {
uint32_t c0 = (uint32_t) counter, c1 = (uint32_t) (counter >> 32);
uint32_t c2 = (uint32_t) seed, c3 = (uint32_t) (seed >> 32);
for (int i = 0; i < 10; ++i) {
const uint32_t hi0 = __umulhi(0x9E3779B9u, c0);
const uint32_t hi1 = __umulhi(0xBB67AE85u, c2);
const uint32_t lo0 = 0x9E3779B9u * c0;
const uint32_t lo1 = 0xBB67AE85u * c2;
const uint32_t n0 = hi1 ^ c1 ^ (uint32_t) i;
const uint32_t n1 = lo1;
const uint32_t n2 = hi0 ^ c3 ^ 0u;
const uint32_t n3 = lo0;
c0 = n0; c1 = n1; c2 = n2; c3 = n3;
}
return (float) (c0 >> 8) * (1.0f / 16777216.0f);
}
// `samp` (device, nullable): [0] 1 = coupled picks, [1] [2] the seed (low, high), [3] top_k, [4] draft temperature (float bits), [5] top_p, [6] min_p, [7] the
// anchor's position. Coupled: row i's pick is the target sampler's inverse-CDF pick over the chain (top_k, top_p on the raw scores, min_p, temperature)
// applied to the 16 scores, with the uniform the verify draws for the row that checks it (Philox(seed, anchor position + i)).
__global__ void select_kernel(const float* __restrict__ hp, const float* __restrict__ unary, const int32_t* __restrict__ tok,
const uint16_t* __restrict__ pred_cb, const uint16_t* __restrict__ succ_cb, int rank, int rows,
const int32_t* __restrict__ anchor_tok, int32_t* __restrict__ out_tok, float* __restrict__ out_prob,
const int32_t* __restrict__ samp) {
const int32_t* __restrict__ anchor_tok, int32_t* __restrict__ out_tok, float* __restrict__ out_prob) {
__shared__ float vp[256];
__shared__ float sc[kDf2TopK];
__shared__ int prev;
@@ -827,63 +804,14 @@ __global__ void select_kernel(const float* __restrict__ hp, const float* __restr
}
__syncthreads();
if (tid == 0) {
if (samp != nullptr && ((const volatile int32_t*) samp)[0] != 0) {
const uint64_t seed = ((uint64_t) (uint32_t) samp[2] << 32) | (uint64_t) (uint32_t) samp[1];
const int top_k = samp[3];
const float temp = __int_as_float(samp[4]), top_p = __int_as_float(samp[5]), min_p = __int_as_float(samp[6]);
const uint64_t counter = (uint64_t) (uint32_t) samp[7] + (uint64_t) i;
// the 16 in the sampler's selection order: score descending, ties to the lower token id
int ord[kDf2TopK];
for (int j = 0; j < kDf2TopK; ++j) ord[j] = j;
for (int a = 1; a < kDf2TopK; ++a) {
const int x = ord[a];
int b = a - 1;
while (b >= 0 && (sc[ord[b]] < sc[x] || (sc[ord[b]] == sc[x] && tok[i * kDf2TopK + ord[b]] > tok[i * kDf2TopK + x]))) { ord[b + 1] = ord[b]; --b; }
ord[b + 1] = x;
}
int k = (top_k > 0 && top_k < kDf2TopK) ? top_k : kDf2TopK;
const float mx = sc[ord[0]];
int n_keep = k;
if (top_p < 1.0f) {
double sum = 0.0;
for (int j = 0; j < k; ++j) sum += exp((double) sc[ord[j]] - (double) mx);
double cum = 0.0;
int cut = k;
for (int j = 0; j < k; ++j) {
cum += exp((double) sc[ord[j]] - (double) mx) / sum;
if (cum >= (double) top_p) { cut = j + 1; break; }
}
n_keep = cut;
}
if (min_p > 0.0f) {
const float thresh = sc[ord[0]] + logf(min_p);
for (int j = 0; j < n_keep; ++j)
if (sc[ord[j]] < thresh) { n_keep = j; break; }
}
const float inv_t = 1.0f / temp;
float smx = sc[ord[0]] * inv_t;
for (int j = 1; j < n_keep; ++j) smx = fmaxf(smx, sc[ord[j]] * inv_t);
double sum = 0.0;
for (int j = 0; j < n_keep; ++j) sum += exp((double) (sc[ord[j]] * inv_t) - (double) smx);
const float u = df2_philox_uniform(seed, counter);
int pick = n_keep - 1;
double cum = 0.0;
for (int j = 0; j < n_keep; ++j) {
cum += exp((double) (sc[ord[j]] * inv_t) - (double) smx) / sum;
if ((double) u < cum) { pick = j; break; }
}
out_tok[i] = tok[i * kDf2TopK + ord[pick]];
out_prob[i] = (float) (exp((double) (sc[ord[pick]] * inv_t) - (double) smx) / sum);
} else {
int bj = 0;
float bv = sc[0];
for (int j = 1; j < kDf2TopK; ++j)
if (sc[j] > bv) { bv = sc[j]; bj = j; }
float z = 0.0f;
for (int j = 0; j < kDf2TopK; ++j) z += __expf(sc[j] - bv);
out_tok[i] = tok[i * kDf2TopK + bj];
out_prob[i] = 1.0f / z;
}
int bj = 0;
float bv = sc[0];
for (int j = 1; j < kDf2TopK; ++j)
if (sc[j] > bv) { bv = sc[j]; bj = j; }
float z = 0.0f;
for (int j = 0; j < kDf2TopK; ++j) z += __expf(sc[j] - bv);
out_tok[i] = tok[i * kDf2TopK + bj];
out_prob[i] = 1.0f / z;
prev = out_tok[i];
}
__syncthreads();
@@ -1085,9 +1013,9 @@ void df2_topk(const float* logits, int n, const int32_t* idx_map, int rows, uint
}
void df2_select(const float* hp, const float* unary, const int32_t* tok, const uint16_t* pred_cb, const uint16_t* succ_cb, int rank, int rows,
const int32_t* anchor_tok, int32_t* out_tok, float* out_prob, void* stream, const int32_t* samp) {
const int32_t* anchor_tok, int32_t* out_tok, float* out_prob, void* stream) {
if (rank != 256) throw std::invalid_argument("df2_select: the selector rank must be 256");
select_kernel<<<1, 256, 0, (cudaStream_t) stream>>>(hp, unary, tok, pred_cb, succ_cb, rank, rows, anchor_tok, out_tok, out_prob, samp);
select_kernel<<<1, 256, 0, (cudaStream_t) stream>>>(hp, unary, tok, pred_cb, succ_cb, rank, rows, anchor_tok, out_tok, out_prob);
check("select");
}
src/program/generate.cpp
-6
View File
@@ -15930,9 +15930,6 @@ int main(int argc, char** argv) {
ver.set_sampling(req_sp);
if (pipe) ver_b.set_sampling(req_sp); // (reaches the later stage's odd verifier)
if (use_mtp) mtp.set_draft_sampling(req_sp); // STRATA_SPEC_COUPLED=1: sampled drafts (a no-op otherwise)
#if defined(STRATA_HAVE_DFLASH2)
if (use_df2) df2.set_sampling(req_sp); // STRATA_DF2_COUPLED=1: the selector's picks drawn with the target's uniform (a no-op otherwise)
#endif
const bool gated_proposal = use_mtp && mtp.prob() && strata::core::spec_gate_pick() &&
gated_proposal_usable(o.pipeline_windows < 2 && o.batch == 0 && n_stages == 1 && !multi_gpu && drive.d.remote_count == 0,
o.lookup_chain, o.adapt_swaps, req_pcie_frac != 0.0 || o.pcie_frac_auto, spec_cost, S,
@@ -18387,9 +18384,6 @@ int main(int argc, char** argv) {
mem_mark("the verifier and the drafter's binding");
ver.set_sampling(sp); // the CLI's own sampling (until 0.1.19 this loop was always greedy); no penalties here
if (use_mtp) mtp.set_draft_sampling(sp); // STRATA_SPEC_COUPLED=1: sampled drafts (a no-op otherwise)
#if defined(STRATA_HAVE_DFLASH2)
if (use_df2) df2.set_sampling(sp); // STRATA_DF2_COUPLED=1: the selector's picks drawn with the target's uniform (a no-op otherwise)
#endif
ver.set_split(o.spec_split);
// auto: the copy kernel for every pack. DMA (the native packs' default until 0.1.13) has the host call
// cudaMemcpyAsync + cudaLaunchHostFunc inside a verify window while the GPU spins on the flag they raise;