Feat/dynamicemb Implementation of LFU_DECAY by costimized JIT function - #437
Feat/dynamicemb Implementation of LFU_DECAY by costimized JIT function#437jiashuy wants to merge 10 commits into
Conversation
Greptile SummaryThis PR implements the
Confidence Score: 4/5The PR is nearly complete and correct, but the Reclaim branch in kernels.cuh leaves a stale LruLfu frequency word when a new key occupies a slot vacated by erase, corrupting LFU eviction ranking for those keys. The Evict path was fixed to clear all score words, but the Reclaim path (slots emptied by table_erase_kernel) was missed. The erase kernel zeroes only word 0; word 1 retains the erased key's accumulated frequency. Any LruLfu table that calls erase() and then re-inserts keys into those slots will silently assign phantom frequency to the new keys, undermining the LFU ordering guarantee. Files Needing Attention: corelib/dynamicemb/src/table_operation/kernels.cuh — the Reclaim branch in the insert() function body needs the same score-block zeroing loop added in the Evict branch. Important Files Changed
Reviews (15): Last reviewed commit: "fix(dynamicemb): make score_function_key..." | Re-trigger Greptile |
| std::snprintf(arch, sizeof(arch), "-arch=sm_%d%d", cc_major, cc_minor); | ||
| const char *opts[] = {arch, "-lto"}; | ||
|
|
||
| nvJitLinkHandle handle; | ||
| auto jl_check = [&](nvJitLinkResult r, const char *what) { | ||
| if (r != NVJITLINK_SUCCESS) { | ||
| std::string log; | ||
| size_t log_size = 0; | ||
| if (nvJitLinkGetErrorLogSize(handle, &log_size) == NVJITLINK_SUCCESS && | ||
| log_size > 0) { | ||
| log.resize(log_size); | ||
| nvJitLinkGetErrorLog(handle, log.data()); | ||
| } | ||
| throw std::runtime_error(std::string("dynamicemb jit_link: ") + what + | ||
| " failed: " + log); | ||
| } | ||
| }; | ||
|
|
||
| if (nvJitLinkCreate(&handle, 2, opts) != NVJITLINK_SUCCESS) | ||
| throw std::runtime_error("dynamicemb jit_link: nvJitLinkCreate failed"); | ||
| jl_check(nvJitLinkAddData(handle, NVJITLINK_INPUT_FATBIN, | ||
| g_custom_fatbin.data(), g_custom_fatbin.size(), | ||
| "evict_custom"), | ||
| "nvJitLinkAddData(custom fatbin)"); | ||
| jl_check(nvJitLinkAddData(handle, NVJITLINK_INPUT_LTOIR, ltoir, ltoir_size, | ||
| "user_score_fn"), | ||
| "nvJitLinkAddData(user ltoir)"); | ||
| jl_check(nvJitLinkComplete(handle), "nvJitLinkComplete"); | ||
| size_t cubin_size = 0; | ||
| jl_check(nvJitLinkGetLinkedCubinSize(handle, &cubin_size), | ||
| "nvJitLinkGetLinkedCubinSize"); | ||
| std::vector<char> cubin(cubin_size); | ||
| jl_check(nvJitLinkGetLinkedCubin(handle, cubin.data()), | ||
| "nvJitLinkGetLinkedCubin"); | ||
| nvJitLinkDestroy(&handle); | ||
|
|
||
| g_custom[key] = load_module(cubin.data(), "cuModuleLoadData(custom cubin)"); | ||
| } | ||
|
|
||
| CUfunction demb_get_evict_fn(int64_t key, bool overflow) { | ||
| std::lock_guard<std::mutex> lk(g_mu); | ||
| if (key == 0) { |
There was a problem hiding this comment.
nvJitLinkHandle leaked when any post-create step throws
nvJitLinkCreate allocates a handle. If any subsequent jl_check(...) call fails, the lambda throws a std::runtime_error and execution jumps past the nvJitLinkDestroy(&handle) call at line 129, leaking the handle. The fix is to wrap the entire block in an RAII guard or use a try/catch that destroys the handle before re-throwing. This includes the failure path of nvJitLinkAddData, nvJitLinkComplete, and the GetLinkedCubin* calls.
90113a2 to
a1bcf2d
Compare
|
/build |
|
❌ Pipeline #57231345 -- failed
Result: 10/17 jobs passed |
| return -float64(cur_timestamp - scores[0]) | ||
|
|
||
|
|
||
| def _lfu_decay(scores, cur_timestamp, gamma): |
There was a problem hiding this comment.
Keep as def score_function(scores, cur_timestamp, cur_frequency)
Verify partial and var_float
a1bcf2d to
2da155f
Compare
|
/build |
|
❌ Pipeline #58387906 -- failed
Result: 11/17 jobs passed |
2da155f to
9ecc190
Compare
|
/build |
|
❌ Pipeline #58407279 -- failed
Result: 13/17 jobs passed |
|
/build |
|
❌ Pipeline #58712240 -- failed
Result: 14/17 jobs passed |
3c7e3c0 to
a2e67b9
Compare
Express LRU+LFU eviction through the compound score strategy (TIMESTAMP, LFU) (either tuple order): word0 = last-access timestamp, word1 = frequency. By default eviction ranks by frequency with an older-timestamp tiebreak; it can be fully customized via a user-supplied score_function. - score_function(scores, cur_timestamp) -> float64 is numba-compiled to LTO-IR and nvJitLink-linked into a prebuilt custom eviction cubin at registration time; the driver launches it for all 2-word eviction. - The default path uses a prebuilt Lex (frequency -> older-timestamp) comparator cubin -- no numba required. - Users index `scores` in logical (tuple) order; subscripts must be integer constants and are statically remapped to physical device words via score_dump_permutation, so (LFU, TIMESTAMP) works transparently. - Custom cubins are cached per (device, function, score order). Two prebuilt fatbins (Lex + custom) are compiled per-arch at build time and shipped as package_data. Tests: 18 table-level tests including whole-set eviction oracles for the default timestamp tiebreak, custom decay, logical!=physical remap, and an end-to-end BatchedDynamicEmbeddingTablesV2 path. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
A full bucket can be evicted by the plain insert path (table_insert), not just insert_and_evict. That path used DefaultEvictor -> the single-score reduce(), which cannot read the 2-word AoS LruLfu layout, so a plain insert that overflowed a bucket would pick a garbage victim (and ignore any custom score_function). Only insert_and_evict was wired to the eviction cubin. Route the num_scores==2 plain-insert eviction through the same cubin: - Factor table_insert_kernel's body into insert_body<Table,KernelTraits,Evictor> so both the AoT kernel (DefaultEvictor) and the cubin (RankedEvictor<Comparator>) share it. - Add a third cubin entry dyn_emb_insert_entry (reuses EvictParams; collects no evicted output) + demb_get_insert_fn(); resolved/cached like the evict entries. - launch_table_insert_kernel routes PolicyType==LruLfu to the cubin (default Lex when score_fn_key==0, else the custom evictor); thread score_fn_key through table_insert and ScoredHashTable.insert(). The plain insert still drops the evicted key (unchanged semantics) but now selects the correct victim (freq->ts tiebreak by default, or the custom decay). Test: table.insert() overflowing a full bucket with _evict_high_freq_first must evict the hot keys (DefaultEvictor would have kept them). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
- compile_evict_fatbins: add --expt-relaxed-constexpr / --expt-extended-lambda so the fatbin TU compiles kernels.cuh (<cub/cub.cuh>, cooperative_groups) with the same flags as the main extension build, avoiding cross-CUDA-version skew. - jit_link.cpp load_from_image: cuModuleUnload the just-loaded module if any cuModuleGetFunction fails, instead of leaking it. - jit_link.cpp link_custom: destroy the nvJitLinkHandle on every exit path via an RAII guard (jl_check / load_from_image can throw), instead of leaking it when linking fails. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
cuda.compile was given only the argument types, so numba inferred the return type from the body. A score_function returning an integer expression (e.g. `return -scores[1]` without a float() cast) would compile to an integer-register return, which the C++ `extern "C" __device__ double user_score_fn(...)` caller reads as junk double bits -- a silent miscompile that degrades eviction to an arbitrary order. Pin the signature to float64(CPointer(uint64), uint64) so numba casts the return value to double instead. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
- evict_lrulfu.cu header: three entry points (evict ovf/noovf + plain insert), reuses insert_body too -- not "two entry points, insert_and_evict_body only". - score_jit.py docstring: register_score_function takes (cc_major, cc_minor). - lru_lfu_score_strategy_design.md: replace the stale "no code written yet" status with an as-built banner (compound tuple instead of an LRU_LFU enum, no gamma, three cubin entries, plain insert also routes through the cubin); keep the body as the original design rationale, code is the source of truth. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
- DynamicEmb_APIs.md: add a "Customized eviction via score_function" section
under DynamicEmbScoreStrategy (signature, logical-order indexing, numba-cuda
requirement, worked example) and a score_function field entry under
DynamicEmbTableOptions. Make the constraint explicit: score_function is
supported ONLY for the two-word {TIMESTAMP, LFU} compound (scores length 2);
any other score_strategy with a score_function raises ValueError.
- example/example.py: add an example lru_lfu_decay_score (time-decayed LFU) and a
--score_strategy {step,timestamp,lru_lfu} flag that wires the compound
(TIMESTAMP, LFU) tuple + score_function into DynamicEmbTableOptions.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Scrub the remaining proposal-era divergences so the design doc describes the shipped implementation (code stays the source of truth), keeping the rationale: - Drop `lfu_decay_gamma` and the `(scores, cur_timestamp, gamma)` signature everywhere; the score_function is `(scores, cur_timestamp) -> float64` with any decay constant written in the body. - No `LRU_LFU` enum: the strategy is the compound `(TIMESTAMP, LFU)` tuple; update the goal, API, example, and component table accordingly. - score_function indexes `scores` in logical (tuple) order with a static logical→physical remap (not "zero remapping"). - Three cubin entries (`dyn_emb_evict_entry_ovf`/`_noovf` + `dyn_emb_insert_entry`) and plain insert also routes through the cubin (two call sites); fix the ABI struct (no gamma field), §11 entry/launch text, and the build-check. - cuda.compile pins the float64 return type; module cache keyed by (device, key). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…e/order - README.md: add a "Customized Eviction Score" section under Usage Notes (compound (TIMESTAMP, LFU) + score_function example, numba-cuda requirement, pointer to the API doc). - DynamicEmb_APIs.md: add a "Compound (multi-score) strategies and score order" section that spells out, in user-facing terms, that a compound strategy keeps more than one score per key and that the tuple order controls two things that always agree — the score column order dumped to file, and how score_function indexes scores — while eviction is identical either way. Reword the score_function / score_strategy text to drop implementation details. - lru_lfu_score_strategy_design.md: fold in the remaining LRU_LFU -> compound-tuple precision fixes (LruLfu spec/strategy/table naming, device-keyed cache) that the prior design-doc rewrite commit missed. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
… guards) Two cheap defensive guards on the JIT eviction path: - UserFnComparator::rank clamps a NaN user score to -inf (evict-first). A NaN compares false against every candidate, so it could never be picked as the eviction minimum, and an all-NaN bucket would find no victim and silently drop the incoming key. Clamping keeps eviction deterministic and inserts progressing. Document the finite-return contract (comparator header + API doc). - reduce()/reduce_ranked() assert num_scores_ == 1 / == 2 respectively, so a misrouted call fails loud instead of silently misreading the layout (nvcc is not built with -DNDEBUG, so the device assert is live). Guards a future regression; after the plain-insert routing both paths already send num_scores==2 to the cubin. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…lti-GPU) score_function_key ignored the device compute capability, so on a single process driving GPUs of different arch the custom score_function registered only the first device's numba LTO-IR (the Python _registered_keys dedup short-circuited the others). A later different-arch device then linked that LTO-IR under the wrong -arch and cuModuleLoadData failed on its first custom-eviction insert. Include cc in the key so each arch gets a distinct key -> its own numba compile and C++ registration (cc is already stored per key, so link_custom picks the right -arch). No C++ change needed. Same-arch devices still share one key/compile with the module cached per device, so homogeneous setups are unaffected. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
9579353 to
f2b242d
Compare
|
/build |
|
❌ Pipeline #60279091 -- failed
Result: 10/14 jobs passed |
Description
Checklist