Skip to content
Closed
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 17 additions & 3 deletions src/program/generate.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -5240,7 +5240,21 @@ int main(int argc, char** argv) {
auto kvg_start = [&](int64_t top) {
kvg.on = strata::core::qsa_kv_elastic() && xcache.vmm_range() != nullptr && d_res != nullptr &&
!host_res.empty() && srcp != nullptr && top > kvg.floor;
if (!kvg.on) return;
if (!kvg.on) {
// The K/V was made elastic at session init, but this run cannot lend it cache slots (no device residency
// table, no slots, a loan at the floor): map the whole window now, from new memory, or the pools keep only
// their first cells and a longer prompt writes past them.
if (strata::core::qsa_kv_elastic() && strata::core::qsa_kv_elastic_cells() < o.max_context) {
if (!strata::core::qsa_kv_elastic_grow(o.max_context, []() -> strata::core::VmmChunk { return 0; })) {
std::fprintf(stderr, "strata generate: the K/V cannot hold %lld cells in VRAM; lower --max-context "
"or start without --kv-grow\n", (long long) o.max_context);
std::exit(1);
}
std::fprintf(stderr, "strata generate: elastic K/V off for this run (no cache slots to lend): the whole "
"window, %lld cells, mapped up front\n", (long long) o.max_context);
}
return;
}
kvg.top = kvg.lo = top;
kvg.cells = strata::core::qsa_kv_elastic_cells();
if (const char* v = std::getenv("STRATA_KV_GROW_STEP"); v != nullptr && std::atoll(v) > 0) kvg.step = std::atoll(v);
Expand Down Expand Up @@ -5332,7 +5346,7 @@ int main(int argc, char** argv) {
if (const strata::core::VmmChunk h = r.unmap(c)) kvg.spare.push_back(h);
}
if (gave > 0 &&
cudaMemcpy(d_res, host_res.data(), host_res.size() * sizeof(int32_t), cudaMemcpyHostToDevice) != cudaSuccess) {
(cudaMemcpy(d_res, host_res.data(), host_res.size() * sizeof(int32_t), cudaMemcpyHostToDevice) != cudaSuccess || cudaStreamSynchronize(cudaStreamLegacy) != cudaSuccess)) {
std::fprintf(stderr, "strata: the residency table upload failed: %s\n", cudaGetErrorString(cudaGetLastError()));
return false;
}
Expand Down Expand Up @@ -5414,7 +5428,7 @@ int main(int argc, char** argv) {
}
for (const auto& [i, s] : filled) host_res[i] = s;
kvg.refilled += (int64_t) filled.size();
if (cudaMemcpy(d_res, host_res.data(), host_res.size() * sizeof(int32_t), cudaMemcpyHostToDevice) != cudaSuccess)
if ((cudaMemcpy(d_res, host_res.data(), host_res.size() * sizeof(int32_t), cudaMemcpyHostToDevice) != cudaSuccess || cudaStreamSynchronize(cudaStreamLegacy) != cudaSuccess))
return false;
for (const strata::core::VmmChunk h : kvg.spare) strata::core::vmm_chunk_free(h); // new ones, if any
kvg.spare.clear();
Expand Down