Perf/fused sparse cache prefetch#430
Conversation
shijieliu
force-pushed
the
perf/fused-sparse-cache-prefetch
branch
from
2a52ef4 to
1d0816a
Compare
Greptile SummaryThis PR replaces the multi-step cache prefetch pipeline (lookup → compact → storage-find → insert-and-evict) with a fused two-kernel design: a cooperative
Confidence Score: 3/5The fused find_or_insert + exchange kernel chain is a substantial rework of the cache prefetch hot path. The most concerning point is that the new exchange kernel uses __trap() to assert that every provisioned cache slot is non-negative; if find_or_insert fails to allocate a slot (all bucket slots in a probing chain are pinned by outstanding batches), the GPU context will terminate with an illegal instruction and no recoverable diagnostic. The PR introduces a hard-crash risk in the exchange kernel via __trap() on negative cache slots, combined with the .item() hot-path synchronization and the always-true non_admitted_mask condition across multiple changed files, warranting extra scrutiny before merging to production. corelib/dynamicemb/src/dynamic_emb_op.cu (exchange kernel __trap), corelib/dynamicemb/src/table_operation/kernels.cuh (unlocked score update), corelib/dynamicemb/dynamicemb/batched_dynamicemb_function.py (non_admitted_mask condition), corelib/dynamicemb/dynamicemb/key_value_table.py (.item() in hot path) Important Files Changed
Sequence Diagram%%{init: {'theme': 'neutral'}}%%
sequenceDiagram
participant PY as Python (prefetch)
participant Cache as DynamicEmbCache
participant FoI as find_or_insert kernel (cooperative)
participant Storage as DynamicEmbStorage
participant Exch as exchange_cache_storage_values kernel
participant Init as initialize_flat kernel
participant Reclaim as reclaim_by_slot kernel
PY->>Cache: find_or_insert(keys, table_ids, scores)
Cache->>FoI: Phase 1 - lookup all keys, pin hits via ref-counter
FoI-->>FoI: cg::this_grid().sync()
FoI->>FoI: Phase 2 - insert misses, evict LRU, write evicted_mask
FoI-->>Cache: indices, founds, evicted_keys/indices/scores/tids, evicted_mask
Cache-->>PY: CacheFindOrInsertResult
PY->>Storage: exchange(CacheExchangeRequest)
Storage->>Storage: lookup(cache_misses) - acquire storage slots
Storage->>Storage: "insert(evicted_keys, publish_and_acquire=True)"
Storage->>Exch: exchange_cache_storage_values(cache-storage rows)
Exch-->>Storage: async value copy complete
Storage-->>PY: CacheExchangeResult(founds, storage_founds, storage_scores)
PY->>PY: admission counter.add / strategy.admit (sparse, skip founds)
PY->>Init: initialize_flat(admitted_mask) - new keys only
PY->>Reclaim: "reclaim(non_admitted_mask) - free rejected slots, set indices=-1"
PY-->>PY: return (slot_indices, slot_indices, non_admitted_mask)
%%{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 Python (prefetch)
participant Cache as DynamicEmbCache
participant FoI as find_or_insert kernel (cooperative)
participant Storage as DynamicEmbStorage
participant Exch as exchange_cache_storage_values kernel
participant Init as initialize_flat kernel
participant Reclaim as reclaim_by_slot kernel
PY->>Cache: find_or_insert(keys, table_ids, scores)
Cache->>FoI: Phase 1 - lookup all keys, pin hits via ref-counter
FoI-->>FoI: cg::this_grid().sync()
FoI->>FoI: Phase 2 - insert misses, evict LRU, write evicted_mask
FoI-->>Cache: indices, founds, evicted_keys/indices/scores/tids, evicted_mask
Cache-->>PY: CacheFindOrInsertResult
PY->>Storage: exchange(CacheExchangeRequest)
Storage->>Storage: lookup(cache_misses) - acquire storage slots
Storage->>Storage: "insert(evicted_keys, publish_and_acquire=True)"
Storage->>Exch: exchange_cache_storage_values(cache-storage rows)
Exch-->>Storage: async value copy complete
Storage-->>PY: CacheExchangeResult(founds, storage_founds, storage_scores)
PY->>PY: admission counter.add / strategy.admit (sparse, skip founds)
PY->>Init: initialize_flat(admitted_mask) - new keys only
PY->>Reclaim: "reclaim(non_admitted_mask) - free rejected slots, set indices=-1"
PY-->>PY: return (slot_indices, slot_indices, non_admitted_mask)
Reviews (1): Last reviewed commit: "perf(dynamicemb): fuse sparse cache-stor..." | Re-trigger Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| non_admitted_in_values[missing_indices] = non_admitted_mask | ||
| initialized_non_admitted = False | ||
| if non_admitted_indices.numel() > 0: | ||
| if non_admitted_mask.numel() > 0: |
There was a problem hiding this comment.
non_admitted_mask.numel() > 0 is always True when any missing keys exist. The intent is to skip the non-admit initializer when no keys were actually rejected, which requires .any() instead.
| if non_admitted_mask.numel() > 0: | |
| if non_admitted_mask.any(): |
| const ExchangeMetadata &metadata, int64_t tile_offset, ValueT *smem) { | ||
| constexpr uint32_t kActiveFlag = | ||
| Direction == kExchangeCacheToStorage | ||
| ? kExchangeMetadataCacheToStorage | ||
| : kExchangeMetadataStorageToCache; | ||
| const bool active = (metadata.flags & kActiveFlag) != 0; |
There was a problem hiding this comment.
__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.
| final_evict_score = *ovf_bucket.scores(local); | ||
| *ovf_bucket.scores(local) = ScoreType(); | ||
| final_evict_index = ovf_iter; | ||
| } else if (ovf_result == InsertResult::Insert) { | ||
| Iter local = ovf_iter - ovf_output_offsets[t_id]; | ||
| table_key_slot = ovf_bucket.keys(local); | ||
| } | ||
| score = Policy::update(ovf_bucket.scores(local), score); | ||
| } | ||
| } |
There was a problem hiding this comment.
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.
| def reclaim( | ||
| self, | ||
| keys: torch.Tensor, | ||
| table_ids: torch.Tensor, | ||
| slot_indices: torch.Tensor, | ||
| mask: torch.Tensor, |
There was a problem hiding this comment.
.item() calls inside the hot prefetch path
DynamicEmbCache.find_or_insert calls int(founds.sum().item()) and int(evicted_mask.sum().item()) whenever self._record_cache_metrics is True. Both .item() calls issue an implicit cudaDeviceSynchronize (or a stream-synchronization event), stalling the CPU until all GPU work up to and including the find_or_insert kernel completes. Because find_or_insert is now in the critical prefetch path (it runs before the storage exchange), this synchronization can negate a significant part of the latency reduction the fused kernel is designed to provide. Consider accumulating metrics asynchronously (e.g., writing them to a pinned/host-visible tensor and reading them only when the user explicitly queries the metrics, after the batch has been consumed by forward).
Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!
1 participant
Description
Checklist