Perf/fused sparse cache prefetch

[shijieliu]

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 Summary

This PR replaces the multi-step cache lookup → value-copy → insert pipeline with a fused cooperative find-or-insert kernel plus a pipelined cp.async exchange kernel, delivering meaningful prefetch throughput improvements for sparse embedding tables.

• Fused find-or-insert: table_insert_and_evict_kernel now launches as a cooperative kernel with a grid-wide barrier separating Phase 1 (lookup/pin hits) from Phase 2 (insert of misses), eliminating a separate kernel launch and host-side synchronization; a rollback path restores evicted slots when backing-storage insert fails.
• Pipelined exchange kernel: exchange_cache_storage_values_kernel uses a 3-stage cp.async SoA pipeline to overlap cache↔storage value transfers bidirectionally; _cache_metrics is moved device-resident to remove a hot-path .item() sync.
• Mask-based admission: non_admitted_positions (int64 indices) is replaced by a non_admitted_mask (bool tensor), propagated through initializer.cu, embedding_admission.py, and batched_dynamicemb_function.py.

Confidence Score: 3/5

Not safe to merge without fixing the counter_flat_index formula and the missing cooperative-launch capability guard; the outstanding_keys_ref exception-safety gap is also a correctness concern.

Three P1 issues: a likely index arithmetic bug in counter_flat_index that could corrupt ref counters for multi-table overflow configurations, a ref-count leak if _scalar_item raises after the increment, and a missing device capability check that can crash on non-cooperative-capable GPUs. These are present defects on the changed code paths.

insert_and_evict.cu (counter_flat_index formula), batched_dynamicemb_function.py (outstanding_keys_ref increment ordering), scored_hashtable.py (cooperative launch capability guard)

Important Files Changed

Filename	Overview
corelib/dynamicemb/src/table_operation/insert_and_evict.cu	Introduces counter_flat_index helper with a likely off-by-base bug for multi-table overflow slots; also adds rollback and reclaim kernels that look correct in isolation.
corelib/dynamicemb/dynamicemb/batched_dynamicemb_function.py	Rewrites _prefetch_cache_path to use fused exchange; outstanding_keys_ref increment is not exception-safe around the blocking _scalar_item call.
corelib/dynamicemb/dynamicemb/scored_hashtable.py	Adds find_or_insert with cooperative kernel launch; missing cudaDevAttrCooperativeLaunch capability guard could crash on unsupported hardware.
corelib/dynamicemb/src/dynamic_emb_op.cu	New exchange_cache_storage_values_kernel with 3-stage cp.async pipeline; contains one dead variable in the 16-byte aligned store branch.
corelib/dynamicemb/src/table_operation/kernels.cuh	InsertKernelTraits extended with FindOrInsert_ template parameter; cooperative two-phase path integrated correctly into table_insert_and_evict_kernel.
corelib/dynamicemb/dynamicemb/key_value_table.py	DynamicEmbCache and DynamicEmbStorage gain find_or_insert/exchange/reclaim methods; _cache_metrics moved to device; _has_live_table_refs guard added.
corelib/dynamicemb/src/initializer.cu	Dispatches on torch.bool dtype for mask-based initialization; CurandStateContext destructor fixed with CUDAGuard + cudaDeviceSynchronize.
corelib/dynamicemb/dynamicemb/types.py	New dataclasses and ABCs for exchange pipeline; DynamicEmbTableState moved here; TYPE_CHECKING guards prevent import cycles.
corelib/dynamicemb/dynamicemb/embedding_admission.py	MultiTableKVCounter and FrequencyAdmissionStrategy updated to accept founds tensor and return bool mask; non_admit_initializer cached correctly.

Sequence Diagram

%%{init: {'theme': 'neutral'}}%%
sequenceDiagram
    participant PY as batched_dynamicemb_function.py
    participant KVT as key_value_table.py
    participant HT as scored_hashtable.py
    participant GPU as CUDA Kernels

    PY->>KVT: cache.find_or_insert(keys)
    KVT->>HT: find_or_insert(keys, scores)
    HT->>GPU: cudaLaunchCooperativeKernel(table_insert_and_evict)
    GPU-->>GPU: Phase 1 - lookup hits, pin cache slots
    GPU-->>GPU: grid sync barrier
    GPU-->>GPU: Phase 2 - insert misses, evict LRU
    GPU-->>HT: hits, miss_slots, evicted_keys
    HT-->>KVT: CacheFindOrInsertResult
    KVT-->>PY: cache_result

    PY->>KVT: storage.exchange(cache_result)
    KVT->>GPU: exchange_cache_storage_values_kernel
    GPU-->>GPU: Stage 0 - async load tiles
    GPU-->>GPU: Stage 1 - prefetch next tiles
    GPU-->>GPU: Stage 2 - store completed tiles
    GPU-->>KVT: CacheExchangeResult

    alt eviction failed
        KVT->>GPU: rollback_failed_evictions_kernel
        GPU-->>KVT: slots restored
        KVT-->>PY: direct_storage_rows fallback
    else eviction succeeded
        KVT-->>PY: cache rows ready
    end

    PY->>KVT: release_cache_exchange_refs()
    KVT->>GPU: decrement ref counters

Loading

%%{init: {'theme': 'base', 'themeVariables': {"darkMode": true, "background": "#0d1117", "primaryColor": "#21262d", "primaryTextColor": "#e6edf3", "primaryBorderColor": "#8b949e", "lineColor": "#8b949e", "textColor": "#e6edf3", "edgeLabelBackground": "#161b22", "actorBkg": "#21262d", "actorBorder": "#8b949e", "actorTextColor": "#e6edf3", "actorLineColor": "#8b949e", "signalColor": "#8b949e", "signalTextColor": "#e6edf3", "noteBkgColor": "#373320", "noteBorderColor": "#d4a72c", "noteTextColor": "#f0e6c0", "labelBoxBkgColor": "#21262d", "labelBoxBorderColor": "#8b949e", "labelTextColor": "#e6edf3", "loopTextColor": "#e6edf3", "activationBkgColor": "#30363d", "activationBorderColor": "#8b949e"}}}%%
sequenceDiagram
    participant PY as batched_dynamicemb_function.py
    participant KVT as key_value_table.py
    participant HT as scored_hashtable.py
    participant GPU as CUDA Kernels

    PY->>KVT: cache.find_or_insert(keys)
    KVT->>HT: find_or_insert(keys, scores)
    HT->>GPU: cudaLaunchCooperativeKernel(table_insert_and_evict)
    GPU-->>GPU: Phase 1 - lookup hits, pin cache slots
    GPU-->>GPU: grid sync barrier
    GPU-->>GPU: Phase 2 - insert misses, evict LRU
    GPU-->>HT: hits, miss_slots, evicted_keys
    HT-->>KVT: CacheFindOrInsertResult
    KVT-->>PY: cache_result

    PY->>KVT: storage.exchange(cache_result)
    KVT->>GPU: exchange_cache_storage_values_kernel
    GPU-->>GPU: Stage 0 - async load tiles
    GPU-->>GPU: Stage 1 - prefetch next tiles
    GPU-->>GPU: Stage 2 - store completed tiles
    GPU-->>KVT: CacheExchangeResult

    alt eviction failed
        KVT->>GPU: rollback_failed_evictions_kernel
        GPU-->>KVT: slots restored
        KVT-->>PY: direct_storage_rows fallback
    else eviction succeeded
        KVT-->>PY: cache rows ready
    end

    PY->>KVT: release_cache_exchange_refs()
    KVT->>GPU: decrement ref counters

Loading

Comments Outside Diff (2)

1. corelib/dynamicemb/dynamicemb/batched_dynamicemb_function.py, line 704-715 (link)

   outstanding_keys_ref incremented before a failable device-to-host sync

   outstanding_keys_ref += num_prefetched_keys is executed before _scalar_item(capacity_check), which performs a blocking .item() device→host copy. If that copy raises (CUDA error, device reset, signal), the ref counter has already been bumped but the capacity rollback branch is never reached, permanently leaking the count and causing subsequent prefetch windows to incorrectly believe more keys are in flight than actually are. The increment should be deferred until after _scalar_item succeeds, or the block must be wrapped in try/finally to decrement on failure.

2. corelib/dynamicemb/dynamicemb/scored_hashtable.py, line 508-545 (link)

   No cudaDevAttrCooperativeLaunch capability check before cooperative kernel launch

   find_or_insert calls cudaLaunchCooperativeKernel (via the cuBin binding), which requires cudaDevAttrCooperativeLaunch == 1 on the target device. The process-level cache _FIND_OR_INSERT_MAX_COOPERATIVE_BLOCKS only stores the block-count ceiling and never verifies cooperative-launch support. On GPUs that do not advertise this attribute (e.g., older Volta/Turing server SKUs in multi-GPU nodes), the kernel launch silently returns cudaErrorInvalidDevice or hangs. A one-time torch.cuda.get_device_properties check for multi_processor_count is already performed; the cooperative-launch attribute should be queried and a clear error raised if unsupported.

_{Reviews (2): Last reviewed commit: "perf(dynamicemb): fused sparse cache-sto..." | Re-trigger Greptile}

[greptile-apps]

__trap() on negative cache row crashes the GPU with no diagnostic

cache_row is cache_slots[input_id], which is the slot_indices tensor produced by find_or_insert. If find_or_insert fails to allocate a slot for any key (e.g., all bucket slots in the relevant probing chain are pinned via positive ref-counts from outstanding batches), slot_indices[i] will be -1. When that value reaches this assertion the kernel issues an illegal instruction (__trap()), which terminates the entire CUDA context with no recoverable diagnostic. The CPU-side outstanding_keys > cache_capacity guard reduces the risk but is a per-batch count, not a per-slot guarantee — with eviction disabled or counters skewed, the check can pass while individual slots still fail to insert. Consider replacing the __trap() with a soft error path (e.g., writing a sentinel to an output error flag tensor checked on the CPU after the kernel) so that a genuinely full cache produces an actionable error rather than a GPU crash.

[greptile-apps]

Score updated without holding the bucket lock

table_update_score_by_slot_kernel writes input_scores[i] directly to *bucket.scores(slot % bucket_capacity) (and the overflow equivalent) without acquiring the per-slot lock. Although the slot's reference counter is positive (the caller owns an acquired reference), concurrent table_lookup_kernel threads that happen to be reading the same slot's score for a non-Const policy path will also hold the lock (via try_lock) while reading, creating a window where a torn 8-byte score read is possible. The score write should acquire the lock with bucket.try_lock before writing and release it afterward, consistent with every other in-place score update in this file.