derive TBE cache associativity from device warp size (#6037)

Summary:

X-link: facebookresearch/FBGEMM#2940

The TBE LRU/LFU/LXU cache hardcoded associativity (ways per set) to 64 on ROCm (`DEFAULT_ASSOC`) and guarded the populate/lookup kernels with `TORCH_CHECK(warp_size() == 64)`, refusing to run on wave32 ROCm devices. Cache associativity must equal the device warp size because one warp cooperatively scans the ways of a single set. The cache kernels are not templated on warp size: device code uses the per-arch `kWarpSize` constant for row indexing, so a multi-arch fat binary already contains correct per-arch device code. Only the host-allocated cache geometry was wrong.

Derive per-instance associativity from the running device (`torch.cuda.get_device_properties(dev).warp_size`) in the training and nbit inference modules, so cache-state tensors match the device kernel's per-arch indexing, and remove the wave64-only `TORCH_CHECK` guards. `DEFAULT_ASSOC` stays as the CPU/no-device fallback. Host-side cache kernel launch configs move to `kWarpSizeHost()`. The cache tests query the device warp size instead of assuming `DEFAULT_ASSOC`, and `lxu_cache_test` now also covers the direct-mapped (associativity 1) path.

Because `DEFAULT_ASSOC` no longer equals the device warp size on ROCm, `cache_test.py`'s `test_lru_cache_insert_large_grid` (which drives the low-level `lru_cache_populate` op directly and must allocate cache geometry matching the kernel's `kWarpSize` row stride) is updated to query the device warp size. Left unchanged it allocates a 32-wide cache on a wave64 device and the kernel, striding by 64, leaves half the slots unpopulated.

One deviation from the combined change (#5804): in `ssd_cache_actions_insert_kernel`, the flat cache-slot index (`cache_set * kWarpSize + insert_slot`) and the conflict-miss stride are `__device__` code and must use the per-arch device constant `kWarpSize`. #5804 changed these to `kWarpSizeHost()`, which calls `at::cuda::warp_size()` -- a host-only function not callable from device code. It is masked on single-arch builds only because `kWarpSizeHost() == kWarpSize` there. These sites are kept as `kWarpSize`; only the host-side launch configs in that file use `kWarpSizeHost()`.

Third in the chain splitting #5804 into reviewable pieces; stacked on #6031.

Authored with assistance from Claude (Anthropic).


Reviewed By: henrylhtsang

Differential Revision: D112940757

Pulled By: q10

Author: jeffdaily

fbgemm_gpu/fbgemm_gpu/split_table_batched_embeddings_ops_inference.py

@@ -394,10 +394,6 @@ def __init__(  # noqa C901
             f"Feature Gates: {[(feature.name, feature.is_enabled()) for feature in FeatureGateName]}"
         )
 
-        # 64 for AMD
-        if cache_assoc == 32 and torch.version.hip is not None:
-            cache_assoc = 64
-
         if device is None:
             self.current_device: torch.device = torch.device(
                 torch.cuda.current_device()
@@ -408,6 +404,20 @@ def __init__(  # noqa C901
             self.current_device = torch.device(device)
         self.use_cpu: bool = self.current_device.type == "cpu"
 
+        # Cache associativity must equal the device warp size (see the training
+        # module's _apply_cache_state for the rationale). Only query a live CUDA
+        # device: on CPU, a meta device (tracing/sharding/publish), or a
+        # CUDA-typed device on a host with no driver, get_device_properties
+        # raises, so fall back to the default associativity in those cases.
+        if (
+            cache_assoc == 32
+            and self.current_device.type == "cuda"
+            and torch.cuda.is_available()
+        ):
+            cache_assoc = torch.cuda.get_device_properties(
+                self.current_device
+            ).warp_size
+
         self.scale_bias_size_in_bytes = scale_bias_size_in_bytes
         self.pooling_mode = pooling_mode
         self.bounds_check_mode_int: int = bounds_check_mode.value

fbgemm_gpu/fbgemm_gpu/split_table_batched_embeddings_ops_training.py

@@ -92,7 +92,9 @@
     pass
 
 
-DEFAULT_ASSOC = 32 if torch.version.hip is None else 64
+# Default cache associativity. Overridden to 64 in each module's
+# __init__ on AMD devices that have 64-lane waves.
+DEFAULT_ASSOC = 32
 INT8_EMB_ROW_DIM_OFFSET = 8
 
 
@@ -3730,6 +3732,20 @@ def _apply_cache_state(
         self.timestep = 1
         self.timesteps_prefetched = []
 
+        # Cache associativity (ways per set) must equal the device warp size:
+        # the cache kernels use one warp to cooperatively scan the ways of a
+        # single set. CDNA (gfx9xx) is warp 64, RDNA (gfx11xx) is warp 32, and
+        # NVIDIA is warp 32. Query the actual device when a live CUDA device is
+        # present, so a single multi-arch ROCm build is correct on both wave
+        # sizes. On CPU, a meta device (tracing/sharding/publish), or a
+        # CUDA-typed device on a host with no driver, get_device_properties
+        # raises, so fall back to the module-level DEFAULT_ASSOC in those cases.
+        self.cache_assoc: int = DEFAULT_ASSOC
+        if self.current_device.type == "cuda" and torch.cuda.is_available():
+            self.cache_assoc = torch.cuda.get_device_properties(
+                self.current_device
+            ).warp_size
+
         self.max_prefetch_depth = MAX_PREFETCH_DEPTH
         self.lxu_cache_locations_list = []
         self.lxu_cache_locations_empty = torch.empty(
@@ -3812,24 +3828,24 @@ def _apply_cache_state(
             assert free_memory > 0
             cache_sets = (
                 int(cache_state.total_cache_hash_size * cache_load_factor)
-                + DEFAULT_ASSOC
+                + self.cache_assoc
                 - 1
-            ) // DEFAULT_ASSOC
+            ) // self.cache_assoc
             cache_sets = 1 if cache_sets == 0 else cache_sets
-            cache_size = cache_sets * DEFAULT_ASSOC * element_size * self.max_D_cache
+            cache_size = cache_sets * self.cache_assoc * element_size * self.max_D_cache
             if cache_size > free_memory:
                 cache_sets = (
                     int(1.0 * free_memory / self.max_D_cache / element_size)
-                    + DEFAULT_ASSOC
+                    + self.cache_assoc
                     - 1
-                ) // DEFAULT_ASSOC
+                ) // self.cache_assoc
         cache_load_factor = (
-            1.0 * cache_sets * DEFAULT_ASSOC / int(cache_state.total_cache_hash_size)
+            1.0 * cache_sets * self.cache_assoc / int(cache_state.total_cache_hash_size)
         )
         assert cache_sets > 0
         if cache_algorithm == CacheAlgorithm.LFU:
             assert cache_sets < 2**24 - 1
-        cache_size = cache_sets * DEFAULT_ASSOC * element_size * self.max_D_cache
+        cache_size = cache_sets * self.cache_assoc * element_size * self.max_D_cache
         self.log(
             f"Using on-device cache with admission algorithm "
             f"{cache_algorithm}, {cache_sets} sets, "
@@ -3862,14 +3878,17 @@ def _apply_cache_state(
         self.register_buffer(
             "lxu_cache_state",
             torch.zeros(
-                cache_sets, DEFAULT_ASSOC, device=self.current_device, dtype=torch.int64
+                cache_sets,
+                self.cache_assoc,
+                device=self.current_device,
+                dtype=torch.int64,
             ).fill_(-1),
         )
         # Cache itself, not auxiliary size
         self.register_buffer(
             "lxu_cache_weights",
             torch.zeros(
-                cache_sets * DEFAULT_ASSOC,
+                cache_sets * self.cache_assoc,
                 self.max_D_cache,
                 device=self.current_device,
                 dtype=dtype,
@@ -3883,7 +3902,7 @@ def _apply_cache_state(
                 size=(
                     (self.total_cache_hash_size + 1,)
                     if cache_algorithm == CacheAlgorithm.LFU
-                    else (cache_sets, DEFAULT_ASSOC)
+                    else (cache_sets, self.cache_assoc)
                 ),
                 device=self.current_device,
                 dtype=torch.int64,
@@ -4028,7 +4047,7 @@ def _init_uvm_cache_counter(self, cache_sets: int, persistent: bool) -> None:
                 "lxu_cache_locking_counter",
                 torch.zeros(
                     cache_sets,
-                    DEFAULT_ASSOC,
+                    self.cache_assoc,
                     device=self.current_device,
                     dtype=torch.int32,
                 ),

fbgemm_gpu/fbgemm_gpu/tbe/cache/cache_config.py

@@ -30,13 +30,12 @@
 import enum
 from typing import NamedTuple
 
-import torch
-
 # Import from config module
 from fbgemm_gpu.tbe.config import EmbeddingLocation
 
-# Default associativity for the UVM cache (32 for CUDA, 64 for ROCm)
-DEFAULT_ASSOC: int = 32 if torch.version.hip is None else 64
+# Default cache associativity. Overridden to 64 in each module's
+# __init__ on AMD devices that have 64-lane waves.
+DEFAULT_ASSOC: int = 32
 
 
 class CacheAlgorithm(enum.Enum):

fbgemm_gpu/src/split_embeddings_cache/lfu_cache_find.cu

@@ -124,9 +124,9 @@ std::pair<Tensor, Tensor> lfu_cache_find_uncached_cuda(
         FBGEMM_LAUNCH_KERNEL(
             (lfu_cache_find_uncached_kernel<index_t>),
             std::min(
-                div_round_up(N, kMaxThreads / kWarpSize),
+                div_round_up(N, kMaxThreads / kWarpSizeHost()),
                 get_max_thread_blocks_for_cache_kernels_()),
-            dim3(kWarpSize, kMaxThreads / kWarpSize),
+            dim3(kWarpSizeHost(), kMaxThreads / kWarpSizeHost()),
             0,
             at::cuda::getCurrentCUDAStream(),
             PTA_B(unique_indices, index_t, 1, 32),

fbgemm_gpu/src/split_embeddings_cache/lfu_cache_populate.cu

@@ -200,9 +200,9 @@ void lfu_cache_insert_cuda(
         FBGEMM_LAUNCH_KERNEL(
             (lfu_cache_insert_kernel<emb_t, cache_t>),
             std::min(
-                div_round_up(N, kCacheMaxThreads / kWarpSize),
+                div_round_up(N, kCacheMaxThreads / kWarpSizeHost()),
                 get_max_thread_blocks_for_cache_kernels_()),
-            dim3(kWarpSize, kCacheMaxThreads / kWarpSize),
+            dim3(kWarpSizeHost(), kCacheMaxThreads / kWarpSizeHost()),
             0,
             at::cuda::getCurrentCUDAStream(),
             PTA_B(weights, emb_t, 1, 64),
@@ -248,15 +248,6 @@ DLL_PUBLIC void lfu_cache_populate_cuda(
 
   CUDA_DEVICE_GUARD(weights);
 
-#ifdef USE_ROCM
-  TORCH_CHECK(
-      at::cuda::warp_size() == 64,
-      __func__,
-      ": TBE cache requires warpSize 64 on ROCm (got ",
-      at::cuda::warp_size(),
-      "); warpSize 32 devices are not yet supported");
-#endif
-
   TORCH_CHECK(
       linear_cache_indices.numel() < std::numeric_limits<int32_t>::max());
   if (linear_cache_indices.numel() == 0) {

fbgemm_gpu/src/split_embeddings_cache/lfu_cache_populate_byte.cu

@@ -176,9 +176,9 @@ void lfu_cache_insert_byte_cuda(
         FBGEMM_LAUNCH_KERNEL(
             (lfu_cache_insert_byte_kernel<index_t>),
             std::min(
-                div_round_up(N, kCacheMaxThreads / kWarpSize),
+                div_round_up(N, kCacheMaxThreads / kWarpSizeHost()),
                 get_max_thread_blocks_for_cache_kernels_()),
-            dim3(kWarpSize, kCacheMaxThreads / kWarpSize),
+            dim3(kWarpSizeHost(), kCacheMaxThreads / kWarpSizeHost()),
             0,
             at::cuda::getCurrentCUDAStream(),
             PTA_B(weights, uint8_t, 1, 64),

fbgemm_gpu/src/split_embeddings_cache/lru_cache_find.cu

@@ -214,15 +214,15 @@ lru_cache_find_uncached_cuda(
         constexpr int PREFETCH_KERNEL_MAX_BLOCKS = 8;
 
         auto grid_size = std::min(
-            div_round_up(N, kMaxThreads / kWarpSize),
+            div_round_up(N, kMaxThreads / kWarpSizeHost()),
             lock_cache_line ? PREFETCH_KERNEL_MAX_BLOCKS
                             : get_max_thread_blocks_for_cache_kernels_());
 
         // Find uncached indices
         FBGEMM_LAUNCH_KERNEL(
             (lru_cache_find_uncached_kernel<index_t>),
             grid_size,
-            dim3(kWarpSize, kMaxThreads / kWarpSize),
+            dim3(kWarpSizeHost(), kMaxThreads / kWarpSizeHost()),
             0,
             at::cuda::getCurrentCUDAStream(),
             PTA_B(unique_indices, index_t, 1, 32),

fbgemm_gpu/src/split_embeddings_cache/lru_cache_populate.cu

@@ -226,7 +226,7 @@ void lru_cache_insert_cuda(
 
         const auto grid_size_uncapped = lock_cache_line
             ? div_round_up(get_device_sm_cnt_(), ALL_TO_PREFETCH_SM_RATIO)
-            : div_round_up(N, kMaxThreads / kWarpSize);
+            : div_round_up(N, kMaxThreads / kWarpSizeHost());
         // HIP enforces a hard limit of 2^32 total threads per launch.
         // lru_cache_insert_kernel grid-strides over n, so capping is
         // correctness-preserving. The lock_cache_line=true branch is already
@@ -240,7 +240,7 @@ void lru_cache_insert_cuda(
         FBGEMM_LAUNCH_KERNEL(
             (lru_cache_insert_kernel<emb_t, cache_t>),
             grid_size,
-            dim3(kWarpSize, kMaxThreads / kWarpSize),
+            dim3(kWarpSizeHost(), kMaxThreads / kWarpSizeHost()),
             0,
             at::cuda::getCurrentCUDAStream(),
             PTA_B(weights, emb_t, 1, 64),
@@ -311,15 +311,6 @@ DLL_PUBLIC void lru_cache_populate_cuda(
 
   CUDA_DEVICE_GUARD(weights);
 
-#ifdef USE_ROCM
-  TORCH_CHECK(
-      at::cuda::warp_size() == 64,
-      __func__,
-      ": TBE cache requires warpSize 64 on ROCm (got ",
-      at::cuda::warp_size(),
-      "); warpSize 32 devices are not yet supported");
-#endif
-
   TORCH_CHECK(
       linear_cache_indices.numel() < std::numeric_limits<int32_t>::max());
   if (linear_cache_indices.numel() == 0) {

fbgemm_gpu/src/split_embeddings_cache/lru_cache_populate_byte.cu

@@ -397,9 +397,9 @@ void lru_cache_insert_byte_cuda(
         FBGEMM_LAUNCH_KERNEL(
             (lru_cache_insert_byte_kernel<index_t>),
             std::min(
-                div_round_up(N, kMaxThreads / kWarpSize),
+                div_round_up(N, kMaxThreads / kWarpSizeHost()),
                 get_max_thread_blocks_for_cache_kernels_()),
-            dim3(kWarpSize, kMaxThreads / kWarpSize),
+            dim3(kWarpSizeHost(), kMaxThreads / kWarpSizeHost()),
             0,
             at::cuda::getCurrentCUDAStream(),
             PTA_B(weights, uint8_t, 1, 64),
@@ -462,9 +462,9 @@ void direct_mapped_lru_cache_insert_byte_cuda(
         FBGEMM_LAUNCH_KERNEL(
             (direct_mapped_lru_cache_insert_byte_kernel<index_t>),
             std::min(
-                div_round_up(N, kMaxThreads / kWarpSize),
+                div_round_up(N, kMaxThreads / kWarpSizeHost()),
                 get_max_thread_blocks_for_cache_kernels_()),
-            dim3(kWarpSize, kMaxThreads / kWarpSize),
+            dim3(kWarpSizeHost(), kMaxThreads / kWarpSizeHost()),
             0,
             at::cuda::getCurrentCUDAStream(),
             PTA_B(weights, uint8_t, 1, 64),

fbgemm_gpu/src/split_embeddings_cache/lxu_cache.cu

@@ -224,8 +224,8 @@ void lxu_cache_locking_counter_decrement_cuda(
 
   auto count = at::zeros_like(lxu_cache_locking_counter);
   const int32_t C = lxu_cache_locking_counter.size(0);
-  TORCH_CHECK(lxu_cache_locking_counter.size(1) == kWarpSize);
-  auto fd = FixedDivisor(kWarpSize);
+  TORCH_CHECK(lxu_cache_locking_counter.size(1) == kWarpSizeHost());
+  auto fd = FixedDivisor(kWarpSizeHost());
 
   const dim3 blocks(
       std::min(
@@ -245,9 +245,9 @@ void lxu_cache_locking_counter_decrement_cuda(
   FBGEMM_LAUNCH_KERNEL(
       lxu_cache_locking_counter_decrement_kernel,
       std::min(
-          div_round_up(C, kMaxThreads / kWarpSize),
+          div_round_up(C, kMaxThreads / kWarpSizeHost()),
           get_max_thread_blocks_for_cache_kernels_()),
-      dim3(kWarpSize, kMaxThreads / kWarpSize),
+      dim3(kWarpSizeHost(), kMaxThreads / kWarpSizeHost()),
       0,
       at::cuda::getCurrentCUDAStream(),
       PTA_B(lxu_cache_locking_counter, int32_t, 2, 32),
@@ -435,19 +435,6 @@ DLL_PUBLIC Tensor lxu_cache_lookup_cuda(
 
   CUDA_DEVICE_GUARD(linear_cache_indices);
 
-#ifdef USE_ROCM
-  // Cache kernels use kWarpSize as the stride for indexing cache rows.
-  // On warpSize 32 ROCm devices (e.g. gfx1100) this produces wrong
-  // addresses because DEFAULT_ASSOC (kWarpSize on host) is 64.
-  // D102579845 introduces kCacheAssoc to fix this; until then, guard.
-  TORCH_CHECK(
-      at::cuda::warp_size() == 64,
-      __func__,
-      ": TBE cache requires warpSize 64 on ROCm (got ",
-      at::cuda::warp_size(),
-      "); warpSize 32 devices are not yet supported");
-#endif
-
   const auto lxu_cache_locations =
       lxu_cache_locations_output.value_or(empty_like(
           linear_cache_indices,
@@ -459,7 +446,7 @@ DLL_PUBLIC Tensor lxu_cache_lookup_cuda(
     return lxu_cache_locations;
   }
 
-  const dim3 threads(kWarpSize, kMaxThreads / kWarpSize);
+  const dim3 threads(kWarpSizeHost(), kMaxThreads / kWarpSizeHost());
   const dim3 blocks(div_round_up(N, kMaxThreads));
 
   AT_DISPATCH_INDEX_TYPES(