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

The final mixer in one fused read

Browse Source 
The hyper-connection read before the head ran gr_read's kernels once per token of the window (lm_head_mix), and once
per draft pass in the draft layer.  Both now go through fused_gr_read_multi, the layers' read (no injection weight: the
final mixer), the window's tokens in one launch.  Its sums are gr_read's: window hashes on fixed text and the
distributions along llama.cpp's texts (scratch/parity_new.sh: IQ3_S and UD-Q4_K_XL, code, thinking, after 16K) equal
the previous build's.

CUPTI, IQ3_S thinking text (--adapt-every 0): the time between windows 3,679 -> 3,578 us a round, each draft step 6-9
us shorter.

Claude-Session: https://claude.ai/code/session_019XxRxQQNyk7Ycdr4ComSs1
This commit is contained in:
Eddoursul committed 2026-09-30 08:53:50 +02:00
1 parent 96af377a19
commit 2ec7a80688
2 files changed
+29 -12
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

src/core/mtp.cpp
+11 -5
View File
@@ -478,11 +478,17 @@ bool MtpDrafter::record_rest(int step_row, cudaStream_t cs, std::string& err) {
gr_write(R_ + (size_t) t * HC * N, y_ + t * N, inj2_ + t * HC, gs, R_ + (size_t) t * HC * N, cs);
}
// ---- the final mixer and the main model's head
for (int t = 0; t < T; ++t)
gr_read(R_ + (size_t) t * HC * N, f32("hyper_connection_mixer.hc_norm.weight"),
bf16("hyper_connection_mixer.input_mix_weight_down.weight"),
bf16("hyper_connection_mixer.input_mix_weight_up.weight"), nullptr, EPS, gs, ss.block.gr,
sample_ + t * N, dummy_inj_, cs);
{
FusedGrArgs fa[kFusedGrMaxT];
for (int t = 0; t < T; ++t) {
fa[t].R = R_ + (size_t) t * HC * N; fa[t].R_out = R_ + (size_t) t * HC * N; fa[t].apply = false;
fa[t].w_norm = f32("hyper_connection_mixer.hc_norm.weight");
fa[t].w_down = bf16("hyper_connection_mixer.input_mix_weight_down.weight");
fa[t].w_up = bf16("hyper_connection_mixer.input_mix_weight_up.weight");
fa[t].eps = EPS; fa[t].lo = lo_ + t * g.hc_lr; fa[t].rs = rs_ + t * HC; fa[t].mixed = sample_ + t * N;
}
fused_gr_read_multi(fa, T, grs_, cs);
}
native_quantize_q8_1(sample_, xq_, (int) N, T, cs);
const bool sub = dhead_ != nullptr;
const int64_t nv = sub ? n_dvocab_ : n_vocab_;
src/core/verify.cpp
+18 -7
View File
@@ -954,15 +954,26 @@ bool Verifier::record_window(int T, cudaStream_t cs, std::string& err) {
const WeightRef *hn = wt.find("output_hc_norm.weight"), *hd = wt.find("output_hc_down.weight"),
*hu = wt.find("output_hc_up.weight");
if (!hn || !hd || !hu) { err = "verify: an output_hc_* weight is missing"; return false; }
for (int t = 0; t < T; ++t) {
BlockBuffers bb = ss.block;
bb.R = Rt(t);
bb.mixed = head_mixed_ + t * N;
if (head_ != nullptr && head_->loaded()) {
if (!lm_head_mix(wt, g, bb, cs, err)) return false;
} else if (!lm_head(wt, g, bb, head_logits_ + (size_t) t * n_vocab_, cs, err)) {
if (head_ != nullptr && head_->loaded()) { // the final mixer, the window's tokens in one read
if (hn->kind != WeightKind::F32 || hd->kind != WeightKind::Bf16InF32 || hu->kind != WeightKind::Bf16InF32) {
err = "verify: the output_hc_* weights have the wrong engine forms";
return false;
}
FusedGrArgs fa[kFusedGrMaxT];
for (int t = 0; t < T; ++t) {
fa[t].R = Rt(t); fa[t].R_out = Rt(t); fa[t].apply = false;
fa[t].w_norm = (const float*) hn->data; fa[t].w_down = (const uint16_t*) hd->data;
fa[t].w_up = (const uint16_t*) hu->data; fa[t].eps = EPS;
fa[t].lo = lo_ + t * g.hc_lr; fa[t].rs = rs_ + t * HC; fa[t].mixed = head_mixed_ + t * N;
}
fused_gr_read_multi(fa, T, grs_, cs);
} else {
for (int t = 0; t < T; ++t) {
BlockBuffers bb = ss.block;
bb.R = Rt(t);
bb.mixed = head_mixed_ + t * N;
if (!lm_head(wt, g, bb, head_logits_ + (size_t) t * n_vocab_, cs, err)) return false;
}
}
if (head_ != nullptr && head_->loaded()) {
try {