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

layer split: auto prices the stage's whole footprint, projections and cudaMalloc granules included

Browse Source 
`ce1e63f` defers a later stage's weight load past the split search so `auto` can trim it, and prices
what that stage will hold from the pack index instead of from `cudaMemGetInfo`.  That price was wrong
in three ways, and where one card is much smaller than the others the small card came out
over-credited, so the search handed it layers it could not hold.  @paulhothersall measured it on a
P40 (24 GB) + RTX 3070 (8 GB) rig: the planner was told the 3070 had 4.26 GiB free where v0.1.39 said
2.34, and it chose K=2 instead of K=16.

- `stage_pool` was re-priced only `if (stage_trim)`, which is false under `auto` - so it held the
  whole pool while the stage really loaded a carved one.
- `NativeDense`'s per-projection `cudaMalloc` was not priced at all: 2018.88 MiB over 300 matrices on
  the iq3_s pack, which is the 1.92 GiB @paulhothersall could not account for.
- `NativeDense::load` makes one `cudaMalloc` per matrix and the driver backs each with 2 MiB granules,
  so a stage's footprint is not the payload sum.  Measured on the RTX 5060: 300 buffers asking
  2018.9 MiB cost 2400.0 MiB - 381.1 MiB of rounding the search could not see.

The carve is now priced per candidate inside `predict`, beside `session_bytes`, and memoised by layer
range: `WeightTable::pool_bytes` for the arena, and a new `NativeDense::weight_bytes_for` for the
projections - the byte walk `load` does, over the GGUF headers, allocating nothing and needing no
device.  `alloc_bytes` prices each matrix the way `cudaMalloc` backs it: a 2 MiB granule for every
matrix of a MiB or more, the payload below.  On this pack's 300 matrices that prices 2255 MiB against
a real 2238, +0.8%, erring high on purpose - the card is left slightly less room than it has rather
than slightly more.  CUDA0 is priced at zero, its arenas being allocated before the search.

Measured on 2 devices (RTX 3060 12 GB + RTX 5060 8 GB), `--layer-split auto`:

- trim off: K=6 -> K=20, which is the split plain v0.1.39 picks with no deferral at all.  The 8 GB
  card's priced footprint goes from 4.06 GiB to 3721 MiB against the 3706 MiB it really takes.  The
  arms that picked K<20 never reached the prompt path at all - "device buffers for a chunk of 512
  tokens do not fit" - so the cost of the wrong split was not a slow run, it was no run.
- trim on: K=6 -> K=17.  Prefill 169.2 -> 364.9 tok/s (+116%), decode 31.2 -> 36.0, expert-cache hit
  41.6% -> 61.5% (300 tokens generated, 10043-token prompt).  The prompt chunk follows the split:
  768 -> 2048 tokens.

On 4 cards (2x RTX 3060 + 2x RTX 5060) the split is unchanged - K=10,19,34 in every arm - and prefill
is 884.0 tok/s against 878.8 for the unfixed build, i.e. noise.  That rig cannot show the flaw: with
four stages no card's cap is the binding constraint.  It is a regression test, and it passes.

A self-check after each stage loads compares the bytes the search priced against the device delta
`cudaMemGetInfo` measures before and after that load, and warns on any under-pricing or on
over-pricing past 5% (the granule rule runs ~0.8% high by design).  It stayed silent in every arm
above, including the two-device runs where the old price was 2.49 GiB out.

Co-Authored-By: Claude Code <[email protected]>
This commit is contained in:
GopiandClaude Code committed 2026-10-06 09:30:54 +05:30
1 parent ce1e63f682
commit 572b607b04
3 files changed
+155 -21
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

include/strata/core/native_dense.hpp
+17
View File
@@ -27,6 +27,16 @@ public:
/// supported type, 2-D), read from the GGUF headers only - so the canonical arena can skip them.
static bool served_names(const std::vector<std::string>& shards, bool include_ple_key,
std::set<std::string>& out, std::string& err);
/// Layer split: the device bytes `load` would allocate for the matrices inside [lb, le) - the same walk, the
/// same filters, the same PLE exception and the same per-matrix granule rounding as `load`, but reading the
/// GGUF headers and the canonical table only, so it allocates nothing and needs no device. A later stage's
/// projections are not in `cudaMemGetInfo` when the split search prices that stage, and the search has to see
/// them: `load` runs one cudaMalloc per matrix, so a search that prices only the canonical arena thinks a
/// small card is emptier than it will be - measured on a P40 (24 GB) + RTX 3070 (8 GB) rig, it chose K=2
/// instead of K=16. `table` must be the canonical table `load` attaches to, and must already hold every
/// tensor the shards carry (`load` looks each one up).
static bool weight_bytes_for(const std::vector<std::string>& shards, WeightTable& table, bool include_ple_key,
int64_t lb, int64_t le, uint64_t& out, std::string& err);
/// Layer split: load only blocks [lb, le) (every other `blk.N.` projection belongs to another GPU's stage; the
/// PLE tensors are loaded everywhere). Process-wide, read by the next `load`; (-1, -1) = all layers.
static void set_layer_range(int lb, int le);
@@ -34,12 +44,19 @@ public:
/// the native kernel also reads, e.g. OrcaRouter's IQ3_XXS) serves the PLE from that row, so it is taken out
/// of `skip` and `load` does not upload the GGUF key over it. A quantized row leaves `skip` unchanged.
static bool keep_unquantized_ple_key(const std::string& pack_dir, std::set<std::string>& skip, std::string& err);
/// The payload of the matrices `load` uploaded - what the weights themselves weigh.
uint64_t weight_bytes() const { return bytes_; }
/// What the device lost to `load`: the same matrices, but each one rounded up to the granule
/// `cudaMalloc` backs it with, because `load` makes one call per matrix. This is the number the layer
/// split has to price; `weight_bytes()` is ~11% smaller on the iq3_s pack and pricing that is what put
/// 42 of 48 layers on an 8 GB card.
uint64_t allocated_bytes() const { return allocated_; }
size_t tensor_count() const { return weights_.size(); }
private:
std::vector<void*> weights_;
void* scratch_ = nullptr;
uint64_t bytes_ = 0;
uint64_t allocated_ = 0;
};
} // namespace strata::core
src/core/native_dense.cpp
+61 -1
View File
@@ -15,6 +15,18 @@
namespace strata::core {
namespace {
// `load` makes one cudaMalloc per projection matrix, and the driver backs each call with 2 MiB granules
// (measured on the RTX 5060: 300 buffers of 6.73 MiB asked 2019 MiB and cost 2400). So a stage's real
// footprint is the granule-rounded sum, not the payload sum. This rule - a granule for every matrix of a
// MiB or more, the payload for the rest - prices the iq3_s pack's 300 matrices at 2255 MiB where they
// really take 2238 (+0.8%). It over-prices, which is the safe side: the card is left slightly less room
// than it has rather than slightly more. Pricing *every* matrix at a granule is far worse (2378 MiB),
// because the driver packs the sub-MiB matrices into the slack of the rounded ones.
constexpr uint64_t kAllocGranule = 2ull << 20;
constexpr uint64_t kSmallAlloc = 1ull << 20;
uint64_t alloc_bytes(uint64_t bytes) {
return bytes < kSmallAlloc ? bytes : (bytes + kAllocGranule - 1) / kAllocGranule * kAllocGranule;
}
int g_layer_lb = -1, g_layer_le = -1; // set_layer_range; -1: every layer
bool in_range(const std::string& name) {
if (g_layer_lb < 0 || name.rfind("blk.", 0) != 0) return true;
@@ -62,6 +74,52 @@ bool NativeDense::served_names(const std::vector<std::string>& shards, bool incl
}
void NativeDense::set_layer_range(int lb, int le) { g_layer_lb = lb; g_layer_le = le; }
// The byte walk `load` does, without the allocations. The filters are `load`'s, in `load`'s order, and the two
// have to stay in step: bytes this prices and `load` does not upload waste cache, bytes `load` uploads and this
// does not price are bytes the split search cannot see. Two deliberate differences, both because `load` has
// already run for CUDA0 by the time the search calls this: a tensor whose canonical ref already carries a native
// override is expected here, and a repeated name is skipped rather than refused.
bool NativeDense::weight_bytes_for(const std::vector<std::string>& shards, WeightTable& table, bool include_ple_key,
int64_t lb, int64_t le, uint64_t& out, std::string& err) {
out = 0;
auto outside = [&](const std::string& name) { // the same rule as `in_range` with (-1,-1) released: [lb, le)
if (name.rfind("blk.", 0) != 0) return false;
const long l = std::strtol(name.c_str() + 4, nullptr, 10);
return l < lb || l >= le;
};
try {
std::set<std::string> seen;
uint64_t total = 0;
for (const auto& path : shards) {
strata::GgufFile gguf(path);
for (const auto& tensor : gguf.tensors()) {
if (!eligible(tensor, include_ple_key)) continue;
if (outside(tensor.name) && tensor.name.find("ple") == std::string::npos) continue;
if (!seen.insert(tensor.name).second) continue;
auto found = table.table_.find(tensor.name);
if (found == table.table_.end()) {
err = "native dense: tensor absent from canonical table: " + tensor.name;
return false;
}
const auto& ref = found->second;
if (!strata::kernels::native_mmvq_supported(tensor.type)) continue;
if (tensor.name == "blk.1.ple_key.weight" && !ref.quantized()) continue;
if (!ref.quantized() || tensor.shape.size() != 2 ||
ref.ne0 <= 0 || ref.ne0 > INT_MAX || ref.ne1 <= 0 || ref.ne1 > INT_MAX ||
tensor.shape[0] != (uint64_t) ref.ne0 || tensor.shape[1] != (uint64_t) ref.ne1) {
err = "native dense: incompatible matrix " + tensor.name;
return false;
}
total += alloc_bytes(strata::kernels::native_mmvq_weight_bytes(tensor.type, (int) ref.ne0, (int) ref.ne1));
}
}
out = total;
return true;
} catch (const std::exception& error) {
err = std::string("native dense: ") + error.what();
return false;
}
}
bool NativeDense::keep_unquantized_ple_key(const std::string& pack_dir, std::set<std::string>& skip,
std::string& err) {
const std::string key = "blk.1.ple_key.weight";
@@ -90,7 +148,7 @@ bool NativeDense::load(const std::vector<std::string>& shards, WeightTable& tabl
std::vector<Pending> pending;
std::set<std::string> seen;
int max_in = 0;
uint64_t total = 0;
uint64_t total = 0, allocated = 0;
uint64_t split_count = 0, split_tensors = 0;
std::set<uint64_t> split_numbers;
bool have_architecture = false;
@@ -183,6 +241,7 @@ bool NativeDense::load(const std::vector<std::string>& shards, WeightTable& tabl
}
max_in = (std::max)(max_in, (int) ref.ne0);
total += bytes;
allocated += alloc_bytes(bytes);
pending.push_back(Pending{&ref, (int) tensor.type, bytes, std::move(data)});
}
}
@@ -201,6 +260,7 @@ bool NativeDense::load(const std::vector<std::string>& shards, WeightTable& tabl
}
scratch_ = scratch.release();
bytes_ = total;
allocated_ = allocated;
return true;
} catch (const std::exception& error) {
err = std::string("native dense: ") + error.what();
src/program/generate.cpp
+77 -20
View File
@@ -2705,24 +2705,54 @@ int main(int argc, char** argv) {
// smaller, which measured cost no decode (K=28 on 9070 XT + R9700: 58.4 tok/s own vs 58.5 borrowing), while a
// search with the reserve moved the boundary to K=32 and decode to 54.8. STRATA_SPLIT_OWN_PLACE=reserve: the
// search sees the reserve.
// ---- WHAT EACH CARD WILL HOLD WHEN THE SEARCH PRICES IT. A later stage's weights load after this search
// (that is what lets `auto` trim them), so its arena is not in `cudaMemGetInfo` yet: size it here from the
// pack index. `pool_bytes` reads the index and returns - it allocates nothing. CUDA0's arena IS up
// already, so its entry is 0, and the number the search sees for every card is the one it has always seen:
// the placement it picks does not move. Deliberate - this buys cache, not a different split.
std::vector<uint64_t> stage_pool(stages.size() + 1, pool_bytes);
stage_pool[0] = 0;
if (stage_trim) {
for (size_t i = 0; i < stages.size(); ++i) {
const int64_t lb = split_at[i], le = i + 1 < split_at.size() ? split_at[i + 1] : g.n_layers;
// ---- WHAT A STAGE HOLDS ONCE ITS WEIGHTS LOAD, PRICED FOR A LAYER RANGE. A later stage's weights load
// after this search (that is what lets `auto` trim them), so they are not in `cudaMemGetInfo` yet and have to
// be priced from the pack instead: `WeightTable::pool_bytes` reads index.txt, `NativeDense::weight_bytes_for`
// reads the GGUF headers, and neither allocates.
// BOTH HAVE TO BE PRICED. `NativeDense::load` runs one cudaMalloc per projection matrix, and the search
// cannot see those - so a search that prices only the canonical arena thinks a card is emptier than it will
// be. That is a small error in GiB and a large one in proportion on a small card: measured by
// @paulhothersall on a P40 (24 GB) + RTX 3070 (8 GB) rig, the planner was told the 3070 had 4.26 GiB free
// where v0.1.39 measured 2.34, and under `auto` it chose K=2 instead of K=16 - 46 layers on the 8 GB card,
// the P40 left on two, -11% prefill at 32K (and -49% with the trim on). The 1.92 GiB it could not see is
// exactly this allocation.
// The carve depends on the layer range - which is what the search is choosing - so it is priced per
// candidate in `predict`, beside `session_bytes`, and memoised by range. CUDA0 is not priced at all: its
// arena and its own native dense are allocated before this search, so they are already missing from `cap[0]`.
bool carve_ok = true;
std::map<std::pair<int64_t, int64_t>, uint64_t> carve_memo;
auto carve_bytes = [&](int64_t lb, int64_t le) -> uint64_t {
const std::pair<int64_t, int64_t> key(lb, le);
if (const auto it = carve_memo.find(key); it != carve_memo.end()) return it->second;
uint64_t total = 0;
if (trim_asked) {
std::set<std::string> sk = skip_base;
add_foreign(lb, le, sk, false);
if (!strata::core::WeightTable::pool_bytes(o.pack, stage_pool[i + 1], err, &sk)) {
std::fprintf(stderr, "strata generate: %s\n", err.c_str());
return 1;
if (!strata::core::WeightTable::pool_bytes(o.pack, total, err, sk.empty() ? nullptr : &sk)) {
carve_ok = false;
return 0;
}
} else {
total = pool_bytes; // not trimming: this stage keeps the whole pool whatever range it runs
}
}
if (!o.native_dense_gguf.empty()) {
uint64_t dense = 0;
if (trim_asked) {
if (!strata::core::NativeDense::weight_bytes_for(o.native_dense_gguf, wt, o.native_ple_key, lb, le,
dense, err)) {
carve_ok = false;
return 0;
}
} else {
// loaded in full above; `set_layer_range` was never called. `allocated_bytes`, not
// `weight_bytes`: what a stage of the SAME shape will really take, granules included.
dense = native_dense.allocated_bytes();
}
total += dense;
}
carve_memo.emplace(key, total);
return total;
};
static const bool place_with_reserve = [] {
const char* v = std::getenv("STRATA_SPLIT_OWN_PLACE");
return v != nullptr && std::string(v) == "reserve";
@@ -2745,8 +2775,7 @@ int main(int argc, char** argv) {
std::vector<double> layer_ms((size_t) ns);
for (int i = 0; i < ns; ++i) {
const int dev = i == 0 ? 0 : stages[(size_t) i - 1]->dev;
cap[(size_t) i] = std::max<int64_t>(stage_room(i == 0 ? -1 : dev, i > 0, i + 1 == ns, true) -
(int64_t) stage_pool[(size_t) i], 0);
cap[(size_t) i] = stage_room(i == 0 ? -1 : dev, i > 0, i + 1 == ns, true);
int sms = 0, khz = 0;
cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, dev);
if (cudaDeviceGetAttribute(&khz, cudaDevAttrClockRate, dev) != cudaSuccess || khz <= 0) khz = 1800000;
@@ -2754,7 +2783,7 @@ int main(int argc, char** argv) {
const double speed = std::max(1.0, (double) sms * (double) khz / 1e6); // SMs x GHz
layer_ms[(size_t) i] = 0.33 * (84.0 * 2.617) / speed;
std::fprintf(stderr, "strata generate: layer split auto: CUDA%d %d SMs at %.2f GHz -> %.2f ms per layer, "
"%.2f GiB free for its weights, session and experts\n", dev, sms, khz / 1e6,
"%.2f GiB free before its weights, session and experts\n", dev, sms, khz / 1e6,
layer_ms[(size_t) i], (double) cap[(size_t) i] / 1073741824.0);
}
const double miss_ms = std::getenv("STRATA_SPLIT_MISS_MS") ? std::atof(std::getenv("STRATA_SPLIT_MISS_MS")) : 190.0;
@@ -2769,11 +2798,15 @@ int main(int argc, char** argv) {
// THE CARVE, PRICED: a placement gives stage i the layers [lb, le), and that range's session is a
// real cost on its device - subtracted here so the search knows what it leaves for experts. This
// is why the sessions are allocated after the search: `session_bytes` is pure arithmetic.
// The range's WEIGHTS are subtracted here too, for the same reason and from the same candidate: a
// later stage's arena and native projections are not allocated yet. CUDA0's are (it is priced at
// zero), which is what makes its `cap` the post-weights number v0.1.39 always printed.
std::vector<int64_t> capr((size_t) ns);
for (int i = 0; i < ns; ++i) {
const int64_t lb = i == 0 ? 0 : at[(size_t) i - 1];
const int64_t le = i + 1 < ns ? at[(size_t) i] : g.n_layers;
capr[(size_t) i] = cap[(size_t) i] - (int64_t) strata::core::session_bytes(g, o.max_context, K, lb, le);
capr[(size_t) i] = cap[(size_t) i] - (int64_t) strata::core::session_bytes(g, o.max_context, K, lb, le) -
(i == 0 ? 0 : (int64_t) carve_bytes(lb, le));
}
std::fill(used.begin(), used.end(), 0);
held_mass = 0;
@@ -2830,6 +2863,10 @@ int main(int argc, char** argv) {
"hold %lld of %zu profiled pairs (~%.1f%% of the routed mass)\n", ks.c_str(), best_ms,
(long long) best_held, profile.size(), 100.0 * best_mass);
}
if (!carve_ok) { // `carve_bytes` reports through `err`; it is called from the search, which cannot return
std::fprintf(stderr, "strata generate: %s\n", err.c_str());
return 1;
}
for (size_t i = 0; i < split_at.size(); ++i)
if (split_at[i] >= g.n_layers) {
std::fprintf(stderr, "strata generate: --layer-split: layer %lld is past the last (%lld)\n",
@@ -2860,8 +2897,8 @@ int main(int argc, char** argv) {
// chosen anything - so a rig that left `--layer-split auto` (the default) kept a full copy of the model on
// every card: on the 4-way IQ3_S rig CUDA1 3650 expert slots, CUDA2 1241, CUDA3 533, and the prompt chunk
// capped at 768 tokens. Carved: 4608 / 2485 / 1772 and a 6144-token chunk. The search above priced these
// arenas from the pack index (`stage_pool`), so the placement it chose is the placement it has always
// chosen; what changed is only where the freed VRAM goes.
// arenas - and the native projections that load beside them - from the pack itself, so what changed here is
// only where the freed VRAM goes.
for (size_t i = 0; i < stages.size(); ++i) {
GpuStage& st = *stages[i];
const strata::core::OnDevice on(st.dev);
@@ -2878,6 +2915,8 @@ int main(int argc, char** argv) {
std::fprintf(stderr, "strata generate: layer split: CUDA%d loads the dense weights of layers %lld-%lld only\n",
st.dev, (long long) lb, (long long) le - 1);
}
size_t took_before = 0, took_total = 0; // what this stage's load really costs, for the check below
cudaMemGetInfo(&took_before, &took_total);
void* arena_s = nullptr;
if (cudaMalloc(&arena_s, pool_s) != cudaSuccess ||
!st.wt.load(o.pack, arena_s, pool_s, err, skip_s.empty() ? nullptr : &skip_s)) {
@@ -2896,6 +2935,24 @@ int main(int argc, char** argv) {
err.c_str());
return 1;
}
// WHAT THE SEARCH WAS TOLD THIS STAGE WOULD HOLD, AGAINST WHAT IT JUST TOOK. The two are the same
// arithmetic by construction (`pool_bytes` on the same skip set; `weight_bytes_for` walks the walk
// `load` walked), so a mismatch means the pricing has drifted from the loading - which is the bug this
// whole block exists to fix, and would otherwise show up only as a wrong split on someone else's rig.
{
size_t took_after = 0, took_total_after = 0;
cudaMemGetInfo(&took_after, &took_total_after);
const uint64_t took = (uint64_t) (took_before - took_after); // arena + projections, measured
const uint64_t priced = carve_bytes(st.lb, st.le);
// UNDER-pricing is the bug this block exists to fix - the search hands a card layers it cannot
// hold, which is how 42 of 48 layers landed on an 8 GB card. A little OVER-pricing is by
// design: the projections are priced a granule at a time (see `alloc_bytes`), which runs ~0.8%
// high. Past that, the pricing has drifted from the loading and someone should hear about it.
if (priced < took || priced - took > took / 20)
std::fprintf(stderr, "strata generate: layer split, CUDA%d: WARNING the split search priced this "
"stage's weights at %llu MiB, it took %llu MiB\n", st.dev,
(unsigned long long) (priced >> 20), (unsigned long long) (took >> 20));
}
size_t fb = 0, tb = 0;
cudaMemGetInfo(&fb, &tb);
std::fprintf(stderr, "strata generate: layer split: CUDA%d holds its weights; %.2f GiB free (its session "