diff --git a/src/kernels/ngram.cpp b/src/kernels/ngram.cpp index 91de3705..1fd7a3a5 100644 --- a/src/kernels/ngram.cpp +++ b/src/kernels/ngram.cpp @@ -508,27 +508,31 @@ bool PleTable::gather_batch(const uint32_t* rows, size_t n_tokens, float* out, s if (impl_->pending) { err = "PleTable::gather_batch while a token is in flight"; return false; } const size_t n = n_tokens * (size_t) PLE_N_HEADS; if (impl_->mode == PleIo::Direct) { - // Fast path: if all requested tokens match our in-flight/completed prefetch slots, collect them directly + // Fast path: match each requested token to any in-flight/completed prefetch slot, then collect directly // without issuing a second PleReader ticket or copying through RowCache. - if (impl_->n_prefetch >= n_tokens && n_tokens <= Impl::kMaxPrefetch) { + if (n_tokens <= Impl::kMaxPrefetch) { + size_t matched[Impl::kMaxPrefetch]; bool exact_match = true; for (size_t t = 0; t < n_tokens; ++t) { - if (std::memcmp(impl_->prefetch_keys[t], rows + t * PLE_N_HEADS, sizeof(uint32_t) * PLE_N_HEADS) != 0) { - exact_match = false; - break; - } + size_t i = 0; + while (i < impl_->n_prefetch && + std::memcmp(impl_->prefetch_keys[i], rows + t * PLE_N_HEADS, + sizeof(uint32_t) * PLE_N_HEADS) != 0) ++i; + if (i == impl_->n_prefetch) { exact_match = false; break; } + matched[t] = i; // duplicate requests may share one immutable prefetch slot } if (exact_match) { + bool collect_ok = true; for (size_t i = 0; i < impl_->n_prefetch; ++i) { - if (!impl_->reader.collect(impl_->prefetch_tickets[i], err)) { - impl_->n_prefetch = 0; - return false; - } + std::string ignored; + if (!impl_->reader.collect(impl_->prefetch_tickets[i], collect_ok ? err : ignored)) + collect_ok = false; // retain the first error, but drain every output buffer } impl_->n_prefetch = 0; + if (!collect_ok) return false; for (size_t t = 0; t < n_tokens; ++t) { for (int h = 0; h < PLE_N_HEADS; ++h) { - impl_->decode(impl_->prefetch_raw[t] + (size_t) h * impl_->rb, + impl_->decode(impl_->prefetch_raw[matched[t]] + (size_t) h * impl_->rb, out + (t * PLE_N_HEADS + (size_t) h) * PLE_HEAD_DIM); } } diff --git a/src/ngram/ple_reader_test.cpp b/src/ngram/ple_reader_test.cpp index 4663455f..670b29cc 100644 --- a/src/ngram/ple_reader_test.cpp +++ b/src/ngram/ple_reader_test.cpp @@ -242,9 +242,89 @@ void bad_sizes_are_refused(const std::string& dir) { } } +// #2622 (reported by @hotschmoe): gather_batch matched the prefetched tokens by position, so a prefetch of A..F +// gathered as D,E,F,A,B,C (or A gathered as A,A,A) issued a second ticket for rows that were already in flight. +// Every case: the bytes equal the mapped reader's, and the requested-row counter says how many rows were re-read. +void prefetched_keys_are_reused(const std::string& dir) { + const k::PleFormatInfo* fmt = nullptr; + for (int i = 0; i < k::ple_format_count(); ++i) + if (std::string(k::ple_formats()[i].name) == "IQ4_NL") fmt = &k::ple_formats()[i]; + if (fmt == nullptr) { CHECK(false, "no IQ4_NL format"); return; } + const std::string path = write_format_table(dir, *fmt, 4096); + if (path.empty()) { CHECK(false, "cannot write the prefetch table"); return; } + struct Case { const char* name; std::vector pre, want; uint64_t extra_rows; }; + std::vector seventeen; + for (int i = 0; i < 17; ++i) seventeen.push_back(i); + const std::vector cases = { + {"ordered", {0, 1}, {0, 1}, 0}, + {"permutation", {0, 1, 2, 3, 4, 5}, {3, 4, 5, 0, 1, 2}, 0}, + {"subset", {0, 1, 2, 3, 4, 5}, {4, 1}, 0}, + {"duplicates", {0}, {0, 0, 0}, 0}, + {"duplicate prefetch", {0, 0, 1}, {0, 1}, 0}, + {"missing key", {0, 1}, {0, 2}, 32}, + {"extra prefetch", {0, 1, 2}, {1}, 0}, + {"none", {0}, {}, 0}, + {"over the window", {0, 1}, seventeen, 17u * k::PLE_N_HEADS}, + {"window overflow", seventeen, {16}, 0}, + }; + const auto keys = [](const std::vector& ids) { + std::vector r; + for (int id : ids) for (int h = 0; h < k::PLE_N_HEADS; ++h) r.push_back((uint32_t) (id * 32 + h)); + return r; + }; + const auto requests = [](const k::PleTable& t) { + unsigned long long n = 0; + const std::string s = t.io_report(); + return std::sscanf(s.c_str(), "ple io: %llu rows", &n) == 1 ? (uint64_t) n : ~0ull; + }; + const size_t tok = (size_t) k::PLE_N_HEADS * k::PLE_HEAD_DIM; + std::vector> refs(cases.size()); + { + k::PleTable mapped; + k::PleIoOptions mio; + mio.mode = k::PleIo::Mmap; + std::string err; + if (!mapped.open(path, err, mio)) { CHECK(false, "mapped open: %s", err.c_str()); return; } + for (size_t c = 0; c < cases.size(); ++c) { + const auto r = keys(cases[c].want); + refs[c].resize(cases[c].want.size() * tok); + CHECK(mapped.gather_batch(r.data(), cases[c].want.size(), refs[c].data(), err), "mapped gather: %s", err.c_str()); + } + } + for (uint64_t cache : {uint64_t(0), uint64_t(256)}) for (bool worker : {false, true}) { + for (size_t c = 0; c < cases.size(); ++c) { + k::PleTable t; + k::PleIoOptions io; + io.mode = k::PleIo::Direct; + io.cache_rows = cache; + io.io_thread = worker; + io.max_inflight = 8; + io.keepalive_ms = 0; + std::string err; + if (!t.open(path, err, io)) { CHECK(false, "direct open: %s", err.c_str()); return; } + const auto pr = keys(cases[c].pre), wr = keys(cases[c].want); + for (size_t i = 0; i < cases[c].pre.size(); ++i) t.prefetch_rows(pr.data() + i * k::PLE_N_HEADS); + const uint64_t rq = requests(t), by = t.bytes_read(); + const float sentinel = -12345.5f; + std::vector out(refs[c].size() + 2, sentinel); + const bool ok = t.gather_batch(wr.data(), cases[c].want.size(), out.data() + 1, err); + CHECK(ok, "%s (cache %llu worker %d): %s", cases[c].name, (unsigned long long) cache, (int) worker, err.c_str()); + CHECK(requests(t) - rq == cases[c].extra_rows, "%s (cache %llu worker %d): %llu rows re-requested, want %llu", + cases[c].name, (unsigned long long) cache, (int) worker, (unsigned long long) (requests(t) - rq), + (unsigned long long) cases[c].extra_rows); + CHECK(t.bytes_read() - by == cases[c].want.size() * k::PLE_N_HEADS * fmt->row_bytes, "%s: bytes_read", cases[c].name); + CHECK(refs[c].empty() || std::memcmp(out.data() + 1, refs[c].data(), refs[c].size() * sizeof(float)) == 0, + "%s (cache %llu worker %d): bytes differ from the mapped reader", cases[c].name, (unsigned long long) cache, (int) worker); + CHECK(out.front() == sentinel && out.back() == sentinel, "%s: wrote outside the output", cases[c].name); + } + } + if (g_fail == 0) std::printf("ple prefetched-key reuse (reordered, duplicated, subset), Direct vs Mmap bytes: OK\n"); +} + void all_formats_round_trip(const std::string& dir) { bf16_widening_is_exact(); // BF16's own claim: a shift, exact for every pattern bad_sizes_are_refused(dir); + prefetched_keys_are_reused(dir); CHECK(k::ple_format_count() > 0, "no formats"); for (int i = 0; i < k::ple_format_count(); ++i) { const k::PleFormatInfo& f = k::ple_formats()[i];