Skip to content

Feat/dynamicemb Implementation of LFU_DECAY by costimized JIT function - #437

Open
jiashuy wants to merge 10 commits into
NVIDIA:mainfrom
jiashuy:feat/dynamicemb-lru-lfu-score-strategy
Open

Feat/dynamicemb Implementation of LFU_DECAY by costimized JIT function#437
jiashuy wants to merge 10 commits into
NVIDIA:mainfrom
jiashuy:feat/dynamicemb-lru-lfu-score-strategy

Conversation

@jiashuy

@jiashuy jiashuy commented Jul 7, 2026

Copy link
Copy Markdown
Collaborator

Description

Checklist

  • I am familiar with the Contributing Guidelines.
  • New or existing tests cover these changes.
  • The documentation is up to date with these changes.

@greptile-apps

greptile-apps Bot commented Jul 7, 2026

Copy link
Copy Markdown
Contributor

Greptile Summary

This PR implements the LFU_DECAY eviction policy for DynamicEmb LruLfu tables by compiling a user-supplied Python decay function via numba into LTO-IR, then linking it into a pre-built CUDA fatbin at runtime using nvJitLink. The new score_function parameter on DynamicEmbTableOptions is optional; omitting it leaves all existing tables on the default lexicographic frequency→timestamp evictor (the Lex fatbin), requiring no numba at runtime.

  • JIT pipeline: score_jit.py validates and remaps the user function (logical→physical score order), pins the return type to float64, includes the device compute capability in the cache key for heterogeneous multi-GPU correctness, and calls demb_register_score_function which builds a per-device CUmodule via link_custom (nvJitLink). RAII guards and try/catch in jit_link.cpp correctly prevent handle and module leaks on all failure paths.
  • Kernel refactor: kernels.cuh extracts insert_body / insert_and_evict_body templated on a pluggable Evictor; evict_lrulfu.cu provides three ABI entry points that instantiate those bodies with RankedEvictor<Comparator>, keeping LruLfu tables on the ranked path and all other tables on the unchanged AoT path.
  • Bug found: The Evict branch inside insert() was correctly updated to clear the full multi-word score block when displacing an existing key, but the Reclaim branch (slot previously vacated by table_erase_kernel) was missed. Word 1 (frequency) retains the erased key's stale value, and ScorePolicy<LruLfu>::update() accumulates on it, giving the incoming key a phantom frequency boost and corrupting LFU eviction order.

Confidence Score: 4/5

The 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

Filename Overview
corelib/dynamicemb/src/table_operation/kernels.cuh Refactors insert/insert_and_evict kernels into pluggable Evictor bodies; adds DefaultEvictor and RankedEvictor. The Evict path correctly clears all score words for LruLfu, but the Reclaim path (reclaimed-from-erase slot) does not, leaving stale frequency word 1 that corrupts LFU ranking for newly-inserted keys.
corelib/dynamicemb/dynamicemb/jit/score_jit.py New JIT glue for LruLfu eviction cubins. Correctly pins float64 return type in numba signature, includes compute capability in the cache key for multi-GPU correctness, and remaps logical→physical score indices before compilation.
corelib/dynamicemb/src/jit/jit_link.cpp New C++ JIT loader: loads Lex fatbin, registers/links custom numba LTO-IR via nvJitLink, and routes per-device CUmodule lookups. RAII HandleGuard and try/catch in load_from_image correctly prevent handle/module leaks on failure paths.
corelib/dynamicemb/src/jit/evict_comparators.cuh Defines LexFreqTsComparator and UserFnComparator. NaN from user_score_fn is correctly clamped to -inf so broken score functions evict deterministically rather than silently dropping keys.
corelib/dynamicemb/src/table_operation/types.cuh Adds assert(num_scores_==1) guard to reduce() and a new reduce_ranked() for 2-word LruLfu layout. LinearBucketTable constructor marked host device for cubin reconstruction.
corelib/dynamicemb/src/jit/evict_lrulfu.cu New translation unit compiled twice (Lex and custom LTO-IR variants). Provides three entry points, reads %globaltimer for cur_ts, and correctly delegates to insert_and_evict_body / insert_body with RankedEvictor.
corelib/dynamicemb/dynamicemb/dynamicemb_config.py Adds score_function field and _score_function_group_key() for table grouping. Validation in post_init correctly rejects score_function with non-compound strategies.
corelib/dynamicemb/dynamicemb/key_value_table.py Wires score_fn_key through create_table_state and _expand_tables_impl so eviction routing is preserved across rehash.
corelib/dynamicemb/test/unit_tests/table_operation/test_lru_lfu.py New test file covering num_scores, frequency accumulation, eviction-by-frequency, gather/scatter, rehash, dump/load, and score_function end-to-end. No test for the Reclaim-path phantom-frequency scenario.
corelib/dynamicemb/setup.py Adds compile_evict_fatbins() to build the two LruLfu fatbins for all DEMB_ARCHS. evict_lrulfu.cu is correctly excluded from the main .so sources.

Reviews (15): Last reviewed commit: "fix(dynamicemb): make score_function_key..." | Re-trigger Greptile

Comment thread corelib/dynamicemb/setup.py
Comment thread corelib/dynamicemb/dynamicemb/batched_dynamicemb_tables.py
Comment on lines +89 to +130
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) {

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 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.

@jiashuy
jiashuy force-pushed the feat/dynamicemb-lru-lfu-score-strategy branch from 90113a2 to a1bcf2d Compare July 8, 2026 07:48
@jiashuy

jiashuy commented Jul 8, 2026

Copy link
Copy Markdown
Collaborator Author

/build

@JacoCheung

JacoCheung commented Jul 8, 2026

Copy link
Copy Markdown
Collaborator

Pipeline #57231345 -- failed

Job Status Log
pre_check ❌ failed view
train_build_x86 ✅ success view
train_build_arm64 ✅ success view
tritonserver_build_x86 ✅ success view
prepare-jet-b200-smoke ✅ success view
prepare-jet-b200-inference ✅ success view
prepare-jet-cw-dfw-e2e-benchmark ✅ success view
build_whl ✅ success view
dynamicemb_test_fwd_bwd_8gpus ❌ failed view
dynamicemb_test_load_dump_8gpus ❌ failed view
unit_test_1gpu_a100 ❌ failed view
unit_test_1gpu_h100 ❌ failed view
unit_test_4gpu ✅ success view
unit_test_tp_4gpu ❌ failed view
L20_unit_test_1gpu ✅ success view
inference_unit_test_1gpu ✅ success view
inference_test_1gpu ❌ failed view

Result: 10/17 jobs passed

View full pipeline

return -float64(cur_timestamp - scores[0])


def _lfu_decay(scores, cur_timestamp, gamma):

@jiashuy jiashuy Jul 10, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Keep as def score_function(scores, cur_timestamp, cur_frequency)
Verify partial and var_float

@jiashuy
jiashuy force-pushed the feat/dynamicemb-lru-lfu-score-strategy branch from a1bcf2d to 2da155f Compare July 17, 2026 02:56
@jiashuy

jiashuy commented Jul 17, 2026

Copy link
Copy Markdown
Collaborator Author

/build

@JacoCheung

JacoCheung commented Jul 17, 2026

Copy link
Copy Markdown
Collaborator

Pipeline #58387906 -- failed

Job Status Log
pre_check ❌ failed view
train_build_x86 ✅ success view
train_build_arm64 ✅ success view
tritonserver_build_x86 ✅ success view
prepare-jet-b200-smoke ✅ success view
prepare-jet-b200-inference ✅ success view
prepare-jet-cw-dfw-e2e-benchmark ✅ success view
build_whl ✅ success view
dynamicemb_test_fwd_bwd_8gpus ❌ failed view
dynamicemb_test_load_dump_8gpus ❌ failed view
unit_test_1gpu_a100 ✅ success view
unit_test_1gpu_h100 ❌ failed view
unit_test_4gpu ✅ success view
unit_test_tp_4gpu ❌ failed view
L20_unit_test_1gpu ✅ success view
inference_unit_test_1gpu ✅ success view
inference_test_1gpu ❌ failed view

Result: 11/17 jobs passed

View full pipeline

Comment thread corelib/dynamicemb/dynamicemb/jit/score_jit.py
@jiashuy
jiashuy force-pushed the feat/dynamicemb-lru-lfu-score-strategy branch from 2da155f to 9ecc190 Compare July 17, 2026 06:14
@jiashuy

jiashuy commented Jul 17, 2026

Copy link
Copy Markdown
Collaborator Author

/build

@JacoCheung

JacoCheung commented Jul 17, 2026

Copy link
Copy Markdown
Collaborator

Pipeline #58407279 -- failed

Job Status Log
pre_check ❌ failed view
train_build_x86 ✅ success view
train_build_arm64 ✅ success view
tritonserver_build_x86 ✅ success view
prepare-jet-b200-smoke ✅ success view
prepare-jet-b200-inference ✅ success view
prepare-jet-cw-dfw-e2e-benchmark ✅ success view
build_whl ✅ success view
dynamicemb_test_fwd_bwd_8gpus ❌ failed view
dynamicemb_test_load_dump_8gpus ✅ success view
unit_test_1gpu_a100 ✅ success view
unit_test_1gpu_h100 ✅ success view
unit_test_4gpu ✅ success view
unit_test_tp_4gpu ❌ failed view
L20_unit_test_1gpu ✅ success view
inference_unit_test_1gpu ✅ success view
inference_test_1gpu ❌ failed view

Result: 13/17 jobs passed

View full pipeline

@jiashuy

jiashuy commented Jul 20, 2026

Copy link
Copy Markdown
Collaborator Author

/build

@JacoCheung

JacoCheung commented Jul 20, 2026

Copy link
Copy Markdown
Collaborator

Pipeline #58712240 -- failed

Job Status Log
pre_check ❌ failed view
train_build_x86 ✅ success view
train_build_arm64 ✅ success view
tritonserver_build_x86 ✅ success view
prepare-jet-b200-smoke ✅ success view
prepare-jet-b200-inference ✅ success view
prepare-jet-cw-dfw-e2e-benchmark ✅ success view
build_whl ✅ success view
dynamicemb_test_fwd_bwd_8gpus ✅ success view
dynamicemb_test_load_dump_8gpus ✅ success view
unit_test_1gpu_a100 ✅ success view
unit_test_1gpu_h100 ✅ success view
unit_test_4gpu ✅ success view
unit_test_tp_4gpu ❌ failed view
L20_unit_test_1gpu ✅ success view
inference_unit_test_1gpu ✅ success view
inference_test_1gpu ❌ failed view

Result: 14/17 jobs passed

View full pipeline

@jiashuy
jiashuy force-pushed the feat/dynamicemb-lru-lfu-score-strategy branch from 3c7e3c0 to a2e67b9 Compare July 20, 2026 07:26
Comment thread corelib/dynamicemb/dynamicemb/jit/score_jit.py Outdated
jiashuy and others added 10 commits July 27, 2026 07:42
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>
@jiashuy
jiashuy force-pushed the feat/dynamicemb-lru-lfu-score-strategy branch from 9579353 to f2b242d Compare July 30, 2026 09:16
@jiashuy

jiashuy commented Jul 30, 2026

Copy link
Copy Markdown
Collaborator Author

/build

@JacoCheung

JacoCheung commented Jul 30, 2026

Copy link
Copy Markdown
Collaborator

Pipeline #60279091 -- failed

Job Status Log
pre_check ❌ failed view
train_build_x86 ✅ success view
train_build_arm64 ❌ failed view
prepare-jet-b200-smoke ⏩ skipped view
prepare-jet-b200-inference ⏩ skipped view
prepare-jet-cw-dfw-e2e-benchmark ✅ success view
build_whl ✅ success view
dynamicemb_test_fwd_bwd_8gpus ✅ success view
dynamicemb_test_load_dump_8gpus ✅ success view
unit_test_1gpu_a100 ✅ success view
unit_test_1gpu_h100 ✅ success view
unit_test_4gpu ✅ success view
unit_test_tp_4gpu ❌ failed view
L20_unit_test_1gpu ✅ success view
inference_unit_test_1gpu ✅ success view
inference_test_1gpu ❌ failed view

Result: 10/14 jobs passed

View full pipeline

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants