diff --git a/corelib/dynamicemb/benchmark/benchmark_selection_kernel_fusion.py b/corelib/dynamicemb/benchmark/benchmark_selection_kernel_fusion.py new file mode 100644 index 000000000..871dba876 --- /dev/null +++ b/corelib/dynamicemb/benchmark/benchmark_selection_kernel_fusion.py @@ -0,0 +1,210 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +""" +Benchmark to verify the fused mask-based erase in table_erase_kernel. + +This script verifies the mask parameter correctly filters which keys are erased +from the admission counter, and that non-admitted keys retain their counter state + +Usage: + torchrun --nproc_per_node=1 benchmark_selection_kernel_fusion.py +""" + +import os +import torch +import torch.distributed as dist +from dynamicemb.batched_dynamicemb_tables import BatchedDynamicEmbeddingTablesV2 +from dynamicemb.embedding_admission import FrequencyAdmissionStrategy, KVCounter +from dynamicemb.dynamicemb_config import ( + DynamicEmbInitializerArgs, + DynamicEmbInitializerMode, + DynamicEmbPoolingMode, + DynamicEmbScoreStrategy, + DynamicEmbTableOptions, +) +from fbgemm_gpu.split_embedding_configs import EmbOptimType +import torchrec + + +def run(): + local_rank = int(os.environ.get("LOCAL_RANK", 0)) + torch.cuda.set_device(local_rank) + dist.init_process_group(backend="nccl") + device = torch.device(f"cuda:{local_rank}") + + emb_dim = 16 + num_embeddings = 10000 + batch_size = 64 + num_iterations = 4 + admission_threshold = 5 + num_features_per_sample = 5 + + value_dim_sgd = emb_dim + total_hbm_bytes = num_embeddings * 4 * value_dim_sgd + local_hbm_bytes = total_hbm_bytes // 2 + + admission_counter_config = KVCounter( + capacity=5000, + bucket_capacity=1024, + key_type=torch.int64, + ) + + admission_strategy = FrequencyAdmissionStrategy( + threshold=admission_threshold, + initializer_args=DynamicEmbInitializerArgs( + mode=DynamicEmbInitializerMode.CONSTANT, + value=0.0, + ), + ) + + table_options = [ + DynamicEmbTableOptions( + embedding_dtype=torch.float32, + dim=emb_dim, + max_capacity=num_embeddings, + local_hbm_for_values=local_hbm_bytes, + index_type=torch.int64, + score_strategy=DynamicEmbScoreStrategy.LFU, + caching=True, + admit_strategy=admission_strategy, + admission_counter=admission_counter_config, + initializer_args=DynamicEmbInitializerArgs( + mode=DynamicEmbInitializerMode.NORMAL, + ), + ) + ] + + model = BatchedDynamicEmbeddingTablesV2( + table_options=table_options, + table_names=["t_0"], + use_index_dedup=True, + pooling_mode=DynamicEmbPoolingMode.NONE, + optimizer=EmbOptimType.SGD, + learning_rate=1e-3, + device=device, + ) + + model.train() + assert model._cache is not None, "Cache must be created" + assert model._admission_counter is not None, "Admission counter must be created" + print(f"Cache: {type(model._cache).__name__}") + print(f"Counter: {type(model._admission_counter).__name__}") + + # Generate keys with controlled frequencies: + # keys 0..49 = frequent (reached admission threshold) + # keys 50..99 = rare (below admission threshold) + torch.manual_seed(42) + total_freq_keys = 40 + total_rare_keys = 80 + keys_per_iter = batch_size * num_features_per_sample + + expected_freq = {} + kjts = [] + for i in range(num_iterations): + g = torch.Generator(device=device) + g.manual_seed(i) + n_freq = int(keys_per_iter * 0.8) + n_rare = keys_per_iter - n_freq + freq_batch = torch.randint( + 0, total_freq_keys, (n_freq,), device=device, generator=g + ) + rare_batch = torch.randint( + total_freq_keys, total_freq_keys + total_rare_keys, + (n_rare,), device=device, generator=g + ) + indices = torch.cat([freq_batch, rare_batch]) + for v in indices: + expected_freq[int(v)] = expected_freq.get(int(v), 0) + 1 + lengths = torch.ones(keys_per_iter, dtype=torch.int64, device=device) + kjts.append(torchrec.KeyedJaggedTensor( + keys=["t_0"], values=indices, lengths=lengths + )) + + # --- Warmup --- + for _ in range(2): + ret = model(kjts[0].values(), kjts[0].offsets()) + grad = torch.empty_like(ret) + ret.backward(grad) + torch.cuda.synchronize() + + # --- Profiled run --- + print("\n=== Profiled forward+backward passes ===") + with torch.profiler.profile( + activities=[torch.profiler.ProfilerActivity.CPU, torch.profiler.ProfilerActivity.CUDA], + record_shapes=True, + with_stack=False, + ) as prof: + for i in range(num_iterations): + ret = model(kjts[i].values(), kjts[i].offsets()) + grad = torch.empty_like(ret) + ret.backward(grad) + torch.cuda.synchronize() + + print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=20)) + + # --- Verify mask-based erase fusion --- + print(f"\n=== Mask-based Erase Fusion Correctness Check ===") + print(f"Threshold: {admission_threshold}, Iterations: {num_iterations}") + + # Dump the admission counter to see which keys are still tracked + counter = model._admission_counter + score_name = counter.score_name_ + counter_keys = {} + for keys, named_scores, _ in counter.table_._batched_export_keys_scores( + [score_name], device, table_id=0 + ): + for k, s in zip(keys, named_scores[score_name]): + counter_keys[int(k)] = int(s) + + admitted_expected = {k for k, v in expected_freq.items() if v >= admission_threshold} + non_admitted_expected = {k for k, v in expected_freq.items() if v < admission_threshold} + + print(f"Unique keys seen: {len(expected_freq)}") + print(f"Expected admitted (freq >= {admission_threshold}): {len(admitted_expected)}") + print(f"Expected non-admitted: {len(non_admitted_expected)}") + print(f"Keys currently in admission counter: {len(counter_keys)}") + + table_keys = set() + cache_state = model._cache._state + for keys, _, _ in cache_state.key_index_map._batched_export_keys_scores( + [cache_state.score_policy.name], device, table_id=0 + ): + table_keys.update(int(x) for x in keys) + for keys, _, _, _ in model._storage.export_keys_values(device, 65536, table_id=0): + table_keys.update(int(x) for x in keys) + missing_from_counter = non_admitted_expected - set(counter_keys.keys()) + in_table = missing_from_counter & table_keys + print(f"Non-admitted keys absent from counter: {len(missing_from_counter)}") + print(f"Of those, found in cache/storage (actually admitted): {len(in_table)}") + truly_missing = missing_from_counter - table_keys + if truly_missing: + print(f"BUG: {len(truly_missing)} keys absent from both counter and storage") + else: + print("All keys accounted for — no unexpected erasures.") + + if truly_missing: + print(f"FAIL: {len(truly_missing)} keys absent from both counter and storage") + print("\n=== FUSION FAILED: Mask-based erase has bugs. ===") + else: + print("\n=== FUSION VERIFIED: Mask correctly skips non-admitted keys. ===") + + # --- Timing --- + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + torch.cuda.synchronize() + n_bench = 30 + start.record() + for _ in range(n_bench): + ret = model(kjts[0].values(), kjts[0].offsets()) + grad = torch.empty_like(ret) + ret.backward(grad) + end.record() + torch.cuda.synchronize() + print(f"\n=== Avg latency ({n_bench} iters): {start.elapsed_time(end) / n_bench:.3f} ms ===") + + dist.barrier() + dist.destroy_process_group() + + +if __name__ == "__main__": + run() diff --git a/corelib/dynamicemb/dynamicemb/batched_dynamicemb_function.py b/corelib/dynamicemb/dynamicemb/batched_dynamicemb_function.py index 68a9d56d1..4e4b36ac2 100644 --- a/corelib/dynamicemb/dynamicemb/batched_dynamicemb_function.py +++ b/corelib/dynamicemb/dynamicemb/batched_dynamicemb_function.py @@ -381,6 +381,7 @@ def _prefetch_cache_path( has_new_in_miss = _bool_item(new_in_miss.any()) if has_new_in_miss and admit_strategy is not None: + new_keys_sub = miss_keys[new_in_miss] new_tids_sub = miss_tids[new_in_miss] @@ -397,7 +398,7 @@ def _prefetch_cache_path( if _bool_item(admit_mask.any()): admission_counter.erase( - new_keys_sub[admit_mask], new_tids_sub[admit_mask] + new_keys_sub, new_tids_sub, mask=admit_mask ) keys_to_insert_mask[new_in_miss] = admit_mask @@ -634,7 +635,7 @@ def _prefetch_hbm_direct_path( admitted_unique_positions = missing_indices[admit_mask] if admit_mask.any(): - admission_counter.erase(admitted_keys, admitted_tids) + admission_counter.erase(missing_keys, missing_table_ids, mask=admit_mask) non_admit = ~admit_mask if non_admit.any(): diff --git a/corelib/dynamicemb/dynamicemb/embedding_admission.py b/corelib/dynamicemb/dynamicemb/embedding_admission.py index fe0f72410..b02595f43 100644 --- a/corelib/dynamicemb/dynamicemb/embedding_admission.py +++ b/corelib/dynamicemb/dynamicemb/embedding_admission.py @@ -91,8 +91,8 @@ def add( self.table_.insert(keys, table_ids, self.score_arg_, score_out=scores_out) return scores_out - def erase(self, keys: torch.Tensor, table_ids: torch.Tensor) -> None: - self.table_.erase(keys, table_ids) + def erase(self, keys: torch.Tensor, table_ids: torch.Tensor, mask=None) -> None: + self.table_.erase(keys, table_ids, mask=mask) def memory_usage(self, mem_type=MemoryType.DEVICE) -> int: return self.table_.memory_usage(mem_type) diff --git a/corelib/dynamicemb/dynamicemb/scored_hashtable.py b/corelib/dynamicemb/dynamicemb/scored_hashtable.py index 021866c9d..fe77ba44c 100644 --- a/corelib/dynamicemb/dynamicemb/scored_hashtable.py +++ b/corelib/dynamicemb/dynamicemb/scored_hashtable.py @@ -168,6 +168,7 @@ def erase( self, keys: torch.Tensor, table_ids: torch.Tensor, + mask: Optional[torch.Tensor] = None ) -> None: """ Erase Keys @@ -787,6 +788,7 @@ def erase( self, keys: torch.Tensor, table_ids: torch.Tensor, + mask: Optional[torch.Tensor] = None, ) -> None: """ Erase Keys @@ -800,6 +802,8 @@ def erase( self.bucket_sizes, keys, table_ids, + indices=None, + mask=mask, ) def load( diff --git a/corelib/dynamicemb/dynamicemb/types.py b/corelib/dynamicemb/dynamicemb/types.py index a8c2b711f..5acacb7de 100644 --- a/corelib/dynamicemb/dynamicemb/types.py +++ b/corelib/dynamicemb/dynamicemb/types.py @@ -337,7 +337,7 @@ def add( ... @abc.abstractmethod - def erase(self, keys: torch.Tensor, table_ids: torch.Tensor) -> None: + def erase(self, keys: torch.Tensor, table_ids: torch.Tensor, mask: Optional[torch.Tensor] = None) -> None: """ Erase keys from the `Counter`. diff --git a/corelib/dynamicemb/src/table_operation/erase.cu b/corelib/dynamicemb/src/table_operation/erase.cu index 33bdcaf1b..55df7073f 100644 --- a/corelib/dynamicemb/src/table_operation/erase.cu +++ b/corelib/dynamicemb/src/table_operation/erase.cu @@ -23,7 +23,7 @@ namespace dyn_emb { void table_erase(at::Tensor table_storage, at::Tensor table_bucket_offsets, int64_t bucket_capacity, at::Tensor bucket_sizes, at::Tensor keys, at::Tensor table_ids, - std::optional indices) { + std::optional indices, std::optional mask) { int64_t num_total = keys.size(0); if (num_total == 0) @@ -32,6 +32,7 @@ void table_erase(at::Tensor table_storage, at::Tensor table_bucket_offsets, auto key_type = get_data_type(keys); auto bucket_sizes_ = get_pointer(bucket_sizes); auto indices_ = get_pointer(indices); + auto mask_ = get_pointer(mask); auto table_ids_ptr = table_ids.data_ptr(); auto table_bucket_offsets_ptr = table_bucket_offsets.data_ptr(); @@ -57,9 +58,9 @@ void table_erase(at::Tensor table_storage, at::Tensor table_bucket_offsets, table_erase_kernel <<<(num_total + BLOCK_SIZE - 1) / BLOCK_SIZE, BLOCK_SIZE, 0, stream>>>( table, table_bucket_offsets_ptr, bucket_sizes_, num_total, keys_, - table_ids_ptr, indices_); + table_ids_ptr, indices_, mask_); }); DEMB_CUDA_KERNEL_LAUNCH_CHECK(); } -} // namespace dyn_emb \ No newline at end of file +} // namespace dyn_emb diff --git a/corelib/dynamicemb/src/table_operation/kernels.cuh b/corelib/dynamicemb/src/table_operation/kernels.cuh index 768bcd0a5..41935da73 100644 --- a/corelib/dynamicemb/src/table_operation/kernels.cuh +++ b/corelib/dynamicemb/src/table_operation/kernels.cuh @@ -552,7 +552,7 @@ __global__ void table_erase_kernel( Table table, int64_t const *__restrict__ table_bucket_offsets, int *__restrict__ bucket_sizes, int64_t batch, typename Table::KeyType const *__restrict__ input_keys, - int64_t const *__restrict__ table_ids, IndexType *__restrict__ indices) { + int64_t const *__restrict__ table_ids, IndexType *__restrict__ indices, bool const *__restrict__ mask) { using KeyType = typename Table::KeyType; using Bucket = typename Table::BucketType; @@ -561,7 +561,13 @@ __global__ void table_erase_kernel( auto tid = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; for (int64_t i = tid; i < batch; i += gridDim.x * blockDim.x) { - + // mask: if non-null, only erase positions where mask[i] is true. + // When mask-skipped, write -1 to indices so callers never read + // uninitialized values from skipped slots. + if (mask && !mask[i]) { + if (indices) indices[i] = -1; + continue; + } KeyType key = input_keys[i]; int64_t hashcode = 0; @@ -787,4 +793,4 @@ __global__ void table_traverse_kernel(Table table, IndexType begin, } } -} // namespace dyn_emb \ No newline at end of file +} // namespace dyn_emb diff --git a/corelib/dynamicemb/src/table_operation/table.cu b/corelib/dynamicemb/src/table_operation/table.cu index 4e4c994d6..b44ca0b82 100644 --- a/corelib/dynamicemb/src/table_operation/table.cu +++ b/corelib/dynamicemb/src/table_operation/table.cu @@ -124,7 +124,8 @@ void bind_table_operation(py::module &m) { m.def("table_erase", &dyn_emb::table_erase, "erase keys from the table", py::arg("table_storage"), py::arg("table_bucket_offsets"), py::arg("bucket_capacity"), py::arg("bucket_sizes"), py::arg("keys"), - py::arg("table_ids"), py::arg("indices") = py::none()); + py::arg("table_ids"), py::arg("indices") = py::none(), + py::arg("mask") = py::none()); m.def("table_export_batch", &dyn_emb::table_export_batch, "export items[offset, offset + batch) from the table", diff --git a/corelib/dynamicemb/src/table_operation/table.cuh b/corelib/dynamicemb/src/table_operation/table.cuh index c42dd494e..f44b17681 100644 --- a/corelib/dynamicemb/src/table_operation/table.cuh +++ b/corelib/dynamicemb/src/table_operation/table.cuh @@ -99,7 +99,8 @@ table_insert_and_evict( void table_erase(at::Tensor table_storage, at::Tensor table_bucket_offsets, int64_t bucket_capacity, at::Tensor bucket_sizes, at::Tensor keys, at::Tensor table_ids, - std::optional indices); + std::optional indices, + std::optional mask); std::tuple table_export_batch(at::Tensor table_storage, int64_t bucket_capacity,