Merge branch 'f-0291-pf-small-df2' into f-0291-pf-small-df2-split

# Conflicts:
#	src/core/arch_defaults_test.cpp
This commit is contained in:
Niko1221 committed 2026-11-02 08:17:43 +01:00
commit dcad09339e
7 files changed
+191 -7

No files matched your search

+27
View File
@@ -122,4 +122,31 @@ inline bool prefill_layer_major_resolve(const char* env, bool hip, int stages, b
return prefill_layer_major_default(hip, stages, v026);
}
/// The small-card prompt path (owner's yes 2026-11-02; DF2-MEM, phaseA-tests/blackwell/RESULTS-DF2MEM.md sections 5-10). On one CUDA card whose expert
/// cache is small, the default prompt path has to read prompts in chunks below 8,192 tokens and streams every expert again for every chunk; reading them
/// with the unfused experts (STRATA_PF_FUSED=0) and, from two chunks on, the layer-major walk (each layer's experts once per super-chunk) is faster:
/// IQ3_S on an RTX PRO 5000 with the cache held to ~1.3K / ~1.6K / ~1.9K engine slots (default chunks 2048 / 3840 / 5888): time to first token of a 20K / 32K
/// prompt -44% / -50%, -41% / -42%, -14% / -12% (and -26% on a real 12 GB RTX 5070 at 20K); from a default chunk of 8,192 on (2.3K slots and up) the default
/// path wins and the same change costs +1..+25%. Not bitwise (first-token KL against the FP16 prompt path 0.012-0.018 for every path including the default;
/// top-1 47 of 48), hence the owner's yes. `default_chunk` is the chunk the DEFAULT path got from the cache sizing (0: none, e.g. no borrowing or an explicit
/// --prefill). One GPU, CUDA, a native (IQ) pack only (what was measured), not under STRATA_DEFAULTS=026. The user's own variables win both ways:
/// `fused_env_set` (STRATA_PF_FUSED named, either value) leaves the expert path alone and with it the walk (a fused layout declines it anyway);
/// `lm_env_set` (STRATA_PREFILL_LAYER_MAJOR named) leaves the walk alone.
struct SmallCardPromptPath {
bool fused_off = false; ///< read the prompt with the unfused experts (STRATA_PF_FUSED=0 as far as the engine is concerned)
bool layer_major = false; ///< and take the layer-major walk for a prompt of two chunks or more (STRATA_PREFILL_LAYER_MAJOR=1)
bool any() const { return fused_off || layer_major; }
};
inline constexpr long long kSmallCardChunkBelow = 8192;
inline SmallCardPromptPath small_card_prompt_path(long long default_chunk, bool hip, int stages, bool native_pack, bool v026, bool fused_env_set,
bool lm_env_set) {
SmallCardPromptPath r;
if (hip || v026 || !native_pack || stages > 1 || default_chunk <= 0 || default_chunk >= kSmallCardChunkBelow) return r;
if (fused_env_set) return r;
r.fused_off = true;
r.layer_major = !lm_env_set;
return r;
}
} // namespace strata
@@ -0,0 +1,18 @@
// include/strata/prefill/pf_fused_override.hpp - one process-wide switch that turns the fused prompt experts off after the environment was read.
//
// STRATA_PF_FUSED is read once by several static caches (moe_fused.cu, moe_fused_iq.cu, prefill.cpp's per-device fused_ring answer). The small-card
// prompt path (strata::small_card_prompt_path in core/sm120_series.hpp) is decided in `serve` only after the expert cache is sized, because it keys on the
// prompt chunk the default path would get, so it cannot set the variable early. It sets this instead; every one of those readers asks it first.
// -1: nothing was decided (the environment and the card's table rule), 0: the fused prompt experts are off for this process.
#pragma once
#include <atomic>
namespace strata::prefill {
inline std::atomic<int>& pf_fused_override() {
static std::atomic<int> v{-1};
return v;
}
} // namespace strata::prefill
+25
View File
@@ -512,6 +512,31 @@ int main() {
CHECK(!prompt_attn_imma_default(false, 60)); // Pascal (P100)
CHECK(!prompt_attn_imma_default(true, 120)); // HIP
}
{
using namespace strata;
// the small-card prompt path: one CUDA card, a native pack, a default prompt chunk below 8,192 tokens
auto sc = [](long long chunk, bool hip = false, int stages = 1, bool native = true, bool v026 = false, bool fe = false, bool le = false) {
return small_card_prompt_path(chunk, hip, stages, native, v026, fe, le);
};
CHECK(sc(2048).fused_off && sc(2048).layer_major); // ~1.3K engine slots (measured -29% / -44% / -50%)
CHECK(sc(3840).fused_off && sc(3840).layer_major);
CHECK(sc(5888).fused_off && sc(5888).layer_major); // ~1.9K slots: still a win (-9..-16%)
CHECK(sc(8191).any()); // the boundary: below 8,192 ...
CHECK(!sc(8192).any()); // ... and not at it (the default is the faster path from there: +7..+25%)
CHECK(!sc(16384).any());
CHECK(!sc(0).any()); // no chunk (no borrowing, or an explicit --prefill): nothing to decide on
CHECK(!sc(-1).any());
CHECK(!sc(2048, true).any()); // a HIP build: no layer-major, the fused experts differ
CHECK(!sc(2048, false, 2).any()); // a layer split: the owner has not decided that one
CHECK(!sc(2048, false, 4).any());
CHECK(sc(2048, false, 0).any()); // (stages 0 = not counted: one card)
CHECK(!sc(2048, false, 1, false).any()); // a non-native (Q2_0) pack: not measured
CHECK(!sc(2048, false, 1, true, true).any()); // STRATA_DEFAULTS=026
CHECK(!sc(2048, false, 1, true, false, true).any()); // STRATA_PF_FUSED named by the user (either value): the user's path
CHECK(sc(2048, false, 1, true, false, false, true).fused_off); // STRATA_PREFILL_LAYER_MAJOR named: the walk is theirs ...
CHECK(!sc(2048, false, 1, true, false, false, true).layer_major); // ... the expert path still follows the rule
CHECK(!sc(2048, false, 1, true, false, true, true).any()); // both named: nothing changes
}
{
using namespace strata;
// the layer-major prompt walk (#3976): off by default on every build (owner's yes 2026-11-02: a layer split reads 32K 37-55% sooner chunk-major on
+3 -2
View File
@@ -18,6 +18,7 @@
// for the k of a half-block: mma k index 16h + 4t + j stands for weight (and activation) 16h + t + 4j. The
// activations are stored in it (perm32), and a lane's two B registers (h = 0 and 1) are adjacent: one 8-byte load.
#include "strata/prefill/moe_fused.hpp"
#include "strata/prefill/pf_fused_override.hpp"
#include "strata/kernels/gfx_arch.hpp"
#include "strata/core/emulate.hpp"
@@ -685,14 +686,14 @@ bool enabled() {
return v == nullptr || v[0] != '0';
#endif
}();
return env && available();
return env && available() && strata::prefill::pf_fused_override().load(std::memory_order_relaxed) != 0;
}
bool requested() {
static const bool env = [] {
const char* v = std::getenv("STRATA_PF_FUSED");
return v != nullptr && v[0] == '1';
}();
return env && available();
return env && available() && strata::prefill::pf_fused_override().load(std::memory_order_relaxed) != 0;
}
size_t act_bytes(int64_t rows, int64_t cols) { return (size_t) rows * (size_t) (cols / 64) * AB; }
+2
View File
@@ -20,6 +20,7 @@
// ldmatrix. The products are mma.sync m16n8k32 (m16n8k16 for the formats with a scale per 16 values); the epilogues
// (SwiGLU, H to int8 per 32 features; down into the per-slot rows) are moe_fused.cu's.
#include "strata/prefill/moe_fused_iq.hpp"
#include "strata/prefill/pf_fused_override.hpp"
#include "strata/kernels/gfx_arch.hpp"
#include "strata/core/emulate.hpp"
#include "strata/core/sm120_series.hpp"
@@ -1999,6 +2000,7 @@ bool native_supported(int gu_type, int d_type) {
return v == nullptr || v[0] == '\0' ? -1 : (v[0] == '1' ? 1 : 0);
}();
if (off || env == 0 || !gu_covered(gu_type)) return false;
if (strata::prefill::pf_fused_override().load(std::memory_order_relaxed) == 0) return false; // the small-card prompt path (pf_fused_override.hpp)
#if !defined(__HIPCC__)
// CUDA: unset, on by default on a compute capability 12.0 card with 80+ SMs (sm120_series.hpp, STRATA_SM120_PERF=0 off) for the gate/up
// formats measured there - the IQ3_S pack's IQ2_S, IQ3_XXS, IQ3_S and IQ4_XS layers with IQ4_NL / Q2_0 down (RTX PRO 5000:
+50 -3
View File
@@ -2,6 +2,7 @@
#include "strata/kernels/fused_gr.hpp"
#include "strata/prefill/prefill.hpp"
#include "strata/prefill/copy_events.hpp"
#include "strata/prefill/pf_fused_override.hpp"
#if defined(STRATA_HAVE_DFLASH2)
#include "strata/kernels/dflash2.hpp"
#endif
@@ -304,6 +305,7 @@ int64_t g_ring_small_max = 0;
// 8192 to 6144, but more experts stay resident). A native pack likewise when the native kernels (moe_fused_iq.hpp)
// take any of its layers: IQ2_XS, 4K / 32K, their first version at 384 slots -3% / -6% against MMQ, at 512 +8% / 0%.
inline bool fused_ring() {
if (strata::prefill::pf_fused_override().load(std::memory_order_relaxed) == 0) return false; // the small-card prompt path decided on MMQ's layout
if (core::peer_portable()) return false; // multi-GPU: --peer-device keeps the MMQ path and its buffer sizes
const strata::kernels::cpu::ExpertLayout& lay = strata::kernels::cpu::expert_layout();
// the Q2_0 pack: fused::enabled(). A native pack: native_supported() alone - it implies enabled() wherever the Q2_0
@@ -6346,9 +6348,9 @@ int64_t Prefill::lm_prepare(int64_t n, std::string& why, int64_t cap) {
return 0;
}
if (embd_rows != nullptr) { why = "an image prompt"; return 0; }
// features of run_impl this walk does not carry: the DFlash2 drafter's prompt taps, and KVA's approximate late layers
// (a prompt that wants them takes the chunk-major path, which has both)
if (!tap_layers_.empty()) { why = "the DFlash2 drafter's prompt taps"; return 0; }
// features of run_impl this walk does not carry: KVA's approximate late layers (a prompt that wants them takes the chunk-major path, which has it).
// The DFlash2 drafter's prompt taps it carries for one stage and one prompt (DF-08: lm_run taps a layer's completed output when the next layer starts)
if (!tap_layers_.empty() && (split || !on_taps || tap_cap_ <= 0)) { why = "the DFlash2 drafter's prompt taps on a layer split"; return 0; }
if (kva_config().on) { why = "STRATA_KVA is on"; return 0; }
if (stop_at_chunk) { why = "a batch read that yields at chunk boundaries"; return 0; }
if (force_chunk_ > 0) { why = "a --prefill-pipe plan sets the chunks of this read"; return 0; } // (left set for run_impl)
@@ -6557,6 +6559,7 @@ bool Prefill::run_multi(const std::vector<Seg>& segs, std::string& err) {
for (const Seg& sg : segs) total += sg.n;
std::string why;
if (hand_in_ != nullptr || !impl_) { err = "prefill: multi-prompt read declined (chunk hand-off in use)"; return false; }
if (!tap_layers_.empty()) { err = "prefill: multi-prompt read declined (the DFlash2 prompt taps are for one prompt)"; return false; }
for (Prefill* q = this; q != nullptr; q = q->next_) q->lm_force_ = true;
struct ForceOff { Prefill* h; ~ForceOff() { for (Prefill* q = h; q != nullptr; q = q->next_) q->lm_force_ = false; } } force_off{this};
if (lm_prepare_chain(total, why) <= 0) { err = "prefill: multi-prompt read declined: " + why; return false; }
@@ -7046,10 +7049,39 @@ bool Prefill::lm_run(const std::vector<Seg>& segs, int depth, const std::vector<
return true;
};
#if defined(STRATA_HAVE_DFLASH2)
// DFlash2 prompt taps (DF-08): one prompt, one stage (lm_prepare declined the rest). A layer's completed output is the rows after the
// hyper-connection write of its FFN half, which this walk does when the NEXT layer starts (or in the final pass for the last one);
// each sub-chunk's tapped rows (positions >= tap_first_) go to their place in the staging, and the super-chunk's rows are handed to
// on_taps, in position order, once all its layers are done
std::vector<int> tap_slot;
const bool tap_on = !tap_layers_.empty() && on_taps && segs.size() == 1 && LB == 0 && LE == L && tap_cap_ > 0;
if (tap_on) {
if (tap_dev_ == nullptr &&
cudaMalloc((void**) &tap_dev_, (size_t) tap_cap_ * tap_layers_.size() * (size_t) g.n_embd * 4) != cudaSuccess) {
(void) cudaGetLastError();
tap_dev_ = nullptr;
err = "prefill: the DFlash2 tap staging does not fit in VRAM";
return false;
}
tap_slot.assign((size_t) g.n_layers, -1);
for (size_t i = 0; i < tap_layers_.size(); ++i) tap_slot[(size_t) tap_layers_[i]] = (int) i;
} else if (!tap_layers_.empty()) {
err = "prefill (layer-major): the DFlash2 taps are not supported for this read";
return false;
}
#endif
for (size_t sci = 0; sci < supers.size(); ++sci) {
const std::vector<LmSub>& subs = supers[sci];
int64_t Sn = 0;
for (const LmSub& sb : subs) Sn = std::max(Sn, sb.t0 + sb.T);
#if defined(STRATA_HAVE_DFLASH2)
// the super-chunk's tapped rows: those at positions >= tap_first_ (the super-chunk's row r is position segs[0].pos0 + subs[0].c0 + r)
const int64_t super_pos = segs[0].pos0 + subs[0].c0;
const int64_t tap_s0 = tap_on ? std::min<int64_t>(Sn, std::max<int64_t>(0, tap_first_ - super_pos)) : Sn; // first tapped row
const int64_t tap_super_rows = Sn - tap_s0;
if (tap_on && tap_super_rows > tap_cap_) { err = "prefill: the DFlash2 tap staging is smaller than the super-chunk's tapped rows"; return false; }
#endif
const int64_t nc = (int64_t) subs.size();
int64_t last_seg0 = -1;
for (int64_t c = 0; c < nc; ++c) if (subs[(size_t) c].seg == 0) last_seg0 = c;
@@ -7147,6 +7179,15 @@ bool Prefill::lm_run(const std::vector<Seg>& segs, int depth, const std::vector<
gr_write(m.R, acc, m.inj, HC, T, cs);
if (strata::kernels::cvec().covers(l - 1))
strata::kernels::cvec_apply(m.R, l - 1, T, D, nullptr, 0, nullptr, 0, false, cs);
#if defined(STRATA_HAVE_DFLASH2)
if (tap_on && tap_slot[(size_t) (l - 1)] >= 0) { // layer l - 1 is complete: its tapped rows of this sub-chunk, mean of the streams
const int64_t lo = std::min<int64_t>(T, std::max<int64_t>(0, tap_s0 - t0));
if (lo < T)
strata::kernels::df2_tap_rows(m.R + lo * D,
tap_dev_ + (size_t) (t0 + lo - tap_s0) * tap_layers_.size() * (size_t) N + (size_t) tap_slot[(size_t) (l - 1)] * (size_t) N,
(int64_t) tap_layers_.size() * N, N, (int) HC, (int) (T - lo), cs);
}
#endif
}
if (l == LE && next_ != nullptr) {
// ---- a layer split: this sub-chunk's rows into the next stage's R_host (on the copy stream, as
@@ -7507,6 +7548,12 @@ bool Prefill::lm_run(const std::vector<Seg>& segs, int depth, const std::vector<
if (l == LE) break;
if (!lm_experts(l, Sn)) return false;
}
#if defined(STRATA_HAVE_DFLASH2)
if (tap_on && tap_super_rows > 0) { // every layer of the super-chunk is done: its tapped rows, in position order (DF-08)
if (cudaStreamSynchronize(cs) != cudaSuccess) { err = std::string("prefill: ") + cudaGetErrorString(cudaGetLastError()); return false; }
if (!on_taps(tap_dev_, tap_super_rows, super_pos + tap_s0, err)) return false;
}
#endif
if (next_ != nullptr) { // the super-chunk's rows are in the next stage's R_host: it reads them on a thread
if (cudaStreamSynchronize(cs) != cudaSuccess || cudaStreamSynchronize(m.copy) != cudaSuccess) {
err = std::string("prefill (layer-major): the layer split's hand-off: ") +
+66 -2
View File
@@ -84,6 +84,8 @@
#include "strata/prefill/kva.hpp"
#include "strata/prefill/kva_floor.hpp"
#include "strata/prefill/prefill.hpp"
#include "strata/prefill/pf_fused_override.hpp"
#include "strata/core/sm120_series.hpp"
#if defined(STRATA_PREFILL_MMQ)
#include "strata/prefill/moe_mmq.hpp"
#endif
@@ -182,6 +184,22 @@ bool window_slot_copy_main(bool into_slot, int slot, int64_t upto, bool kv_only,
#include <unordered_map>
#include <vector>
// STRATA_DF2_TAP_DUMP=FILE (research, off by default): every prompt tap batch the DFlash2 context takes is appended to FILE as
// int64 pos0, int64 rows, then rows x floats_per_row float32 (the tap rows [rows][taps][n_embd]), for the layer-major / chunk-major
// equivalence test (DF-08). Nothing runs without the variable.
static void df2_tap_dump(const float* dev, int64_t rows, int64_t pos0, int64_t floats_per_row) {
static const char* path = std::getenv("STRATA_DF2_TAP_DUMP");
if (path == nullptr || path[0] == '\0' || rows <= 0) return;
std::vector<float> h((size_t) (rows * floats_per_row));
if (cudaMemcpy(h.data(), dev, h.size() * sizeof(float), cudaMemcpyDeviceToHost) != cudaSuccess) return;
if (std::FILE* f = std::fopen(path, "ab")) {
std::fwrite(&pos0, 8, 1, f);
std::fwrite(&rows, 8, 1, f);
std::fwrite(h.data(), sizeof(float), h.size(), f);
std::fclose(f);
}
}
// logprobs=K / allowed=...: one `LP` line for the token just emitted, from `ver`'s head
// logits for the window row it came from (before sampling, penalties and temperature).
// Format matches the SYCL port's LP line so one server-side reader serves both.
@@ -9904,6 +9922,52 @@ int main(int argc, char** argv) {
return std::max(c, small); // never a smaller chunk than 0.1.39's
};
int64_t chunk = scan(nullptr);
// The small-card prompt path (strata::small_card_prompt_path; owner's yes 2026-11-02): this chunk is what the DEFAULT prompt path got from the cache.
// Below 8,192 tokens the default streams every expert again for every chunk, and the unfused experts plus (from two chunks) the layer-major walk
// read the same prompt 9-50% sooner (RESULTS-DF2MEM.md); the chunk is then sized again for that layout. One CUDA card, a native pack, an auto
// chunk; STRATA_PF_FUSED / STRATA_PREFILL_LAYER_MAJOR set by the user win; STRATA_PF_SMALL=0 turns the rule off.
if (chunk > 0 && o.prefill_auto && !multi_gpu && stages.empty() && !(std::getenv("STRATA_PF_SMALL") != nullptr && std::getenv("STRATA_PF_SMALL")[0] == '0')) {
#if defined(STRATA_USE_HIP)
constexpr bool pf_small_hip = true;
#else
constexpr bool pf_small_hip = false;
#endif
const char* fe = std::getenv("STRATA_PF_FUSED");
const char* le = std::getenv("STRATA_PREFILL_LAYER_MAJOR");
const strata::SmallCardPromptPath pp = strata::small_card_prompt_path(
chunk, pf_small_hip, 1, strata::kernels::cpu::expert_layout().native, strata::qwen_defaults_026(),
fe != nullptr && fe[0] != '\0', le != nullptr && le[0] != '\0');
if (pp.any()) {
const int64_t chunk_default = chunk;
if (pp.fused_off) strata::prefill::pf_fused_override().store(0);
if (pp.layer_major) {
#if defined(_WIN32)
_putenv_s("STRATA_PREFILL_LAYER_MAJOR", "1");
#else
setenv("STRATA_PREFILL_LAYER_MAJOR", "1", 0);
#endif
}
chunk = scan(nullptr);
if (chunk <= 0) { // the unfused layout's buffers do not fit anywhere: keep the default path as it was
strata::prefill::pf_fused_override().store(-1);
if (pp.layer_major) {
#if defined(_WIN32)
_putenv_s("STRATA_PREFILL_LAYER_MAJOR", "");
#else
unsetenv("STRATA_PREFILL_LAYER_MAJOR");
#endif
}
chunk = scan(nullptr);
std::fprintf(stderr, "strata serve: small expert cache: the unfused prompt path does not fit here; the default path stays (chunk %lld tokens)\n",
(long long) chunk);
} else
std::fprintf(stderr, "strata serve: small expert cache: the default prompt path would read in chunks of %lld tokens (under %lld), "
"so prompts read with the unfused experts%s (chunk %lld tokens); STRATA_PF_SMALL=0 keeps the default path, "
"STRATA_PF_FUSED / STRATA_PREFILL_LAYER_MAJOR set by hand win\n",
(long long) chunk_default, (long long) strata::kSmallCardChunkBelow,
pp.layer_major ? " and, from two chunks, the layer-major walk" : "", (long long) chunk);
}
}
// #448: a small card in a layer split caps every stage's chunk (an RTX 3080's 512-slot cache held a
// 32 GB card's split to 512 tokens: prompts 6.2x slower, decode the same). Named when it bites, so the
// regression is one log line: each stage that cannot fund the chunk CUDA0 alone would read in.
@@ -10645,7 +10709,7 @@ int main(int argc, char** argv) {
"a draft head 0.2 GiB, 0.1 GiB of prompt staging); --vram-reserve-mib 1900 keeps it. The MTP drafts instead.\n", why_not.c_str());
use_df2 = false;
} else {
sp.on_taps = [&](const float* taps, int64_t rows, int64_t pos0, std::string& e) -> bool { return df2.context(taps, rows, pos0, e); };
sp.on_taps = [&](const float* taps, int64_t rows, int64_t pos0, std::string& e) -> bool { df2_tap_dump(taps, rows, pos0, (int64_t) df2.tap_layers().size() * g.n_embd); return df2.context(taps, rows, pos0, e); };
std::fprintf(stderr, "strata serve: DFlash2 drafter on (%s, graphs captured in %.0f ms); windows up to %d tokens; the MTP layer stays as the fallback\n",
o.draft_model.c_str(), df2.ms_capture, o.spec);
}
@@ -18571,7 +18635,7 @@ int main(int argc, char** argv) {
if (use_df2) { // the drafter's context: the last 2047 prompt rows' taps
df2.reset();
prefill.set_taps(df2.tap_layers(), std::max<int64_t>(0, n_prompt - 2047), 2048);
prefill.on_taps = [&](const float* taps, int64_t rows, int64_t pos0, std::string& e) -> bool { return df2.context(taps, rows, pos0, e); };
prefill.on_taps = [&](const float* taps, int64_t rows, int64_t pos0, std::string& e) -> bool { df2_tap_dump(taps, rows, pos0, (int64_t) df2.tap_layers().size() * g.n_embd); return df2.context(taps, rows, pos0, e); };
}
#endif
if (!kvg_started) { // the elastic K/V: every cell this run can reach (no loan)