Skip to content

Add grid-stride + ROCm cap to linearize_cache_indices kernels #6098

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
q10 wants to merge 2 commits into
base: main
Choose a base branch
from
+117 −75
Open

Add grid-stride + ROCm cap to linearize_cache_indices kernels #6098

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
6d93817
Add grid-stride + ROCm cap to lxu_cache_lookup_kernel
q10 Jul 30, 2026
b1bb8c8
Add grid-stride + ROCm cap to linearize_cache_indices kernels
q10 Jul 30, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
97 changes: 60 additions & 37 deletions fbgemm_gpu/src/split_embeddings_cache/linearize_cache_indices.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@
*/

#include "common.cuh"
#include "fbgemm_gpu/utils/cuda_utilities.cuh"

using Tensor = at::Tensor;
using namespace fbgemm_gpu;
Expand All @@ -24,36 +25,45 @@ __global__ __launch_bounds__(kMaxThreads) void linearize_cache_indices_kernel(
pta::PackedTensorAccessor32<int64_t, 1, at::RestrictPtrTraits>
linear_cache_indices,
const int64_t indices_base_offset) {
#ifdef USE_ROCM
for (index_t index = blockIdx.x * blockDim.x + threadIdx.x;
index < indices.size(0);
index += blockDim.x * gridDim.x) {
#else
const index_t index = blockIdx.x * blockDim.x + threadIdx.x;
if (index >= indices.size(0)) {
return;
}
#endif

// Perform binary search.
int left = 0;
int right = table_offsets.size(0);
const auto index_with_offset = index + indices_base_offset;
while (left != right) {
const int middle =
left + (right - left) / 2; // Avoid overflow in midpoint calculation
if (table_offsets[middle] <= index_with_offset) {
left = middle + 1;
} else {
right = middle;
}
}
const int table_index = left;

// Perform binary search.
int left = 0;
int right = table_offsets.size(0);
const auto index_with_offset = index + indices_base_offset;
while (left != right) {
const int middle =
left + (right - left) / 2; // Avoid overflow in midpoint calculation
if (table_offsets[middle] <= index_with_offset) {
left = middle + 1;
const auto max_offset =
::__ldg(&cache_hash_size_cumsum[cache_hash_size_cumsum.size(0) - 1]);
const auto curr_offset = ::__ldg(&cache_hash_size_cumsum[table_index]);
if (curr_offset >= 0 && indices[index] >= 0) {
linear_cache_indices[index] = indices[index] + curr_offset;
} else {
right = middle;
// Either table index is wrong, or index value is negative (due to
// pruning): set it to invalid value.
linear_cache_indices[index] = max_offset;
}
}
const int table_index = left;

const auto max_offset =
::__ldg(&cache_hash_size_cumsum[cache_hash_size_cumsum.size(0) - 1]);
const auto curr_offset = ::__ldg(&cache_hash_size_cumsum[table_index]);
if (curr_offset >= 0 && indices[index] >= 0) {
linear_cache_indices[index] = indices[index] + curr_offset;
} else {
// Either table index is wrong, or index value is negative (due to pruning):
// set it to invalid value.
linear_cache_indices[index] = max_offset;
}
#ifdef USE_ROCM
} // for index (grid-stride loop, ROCm only)
#endif
}

} // namespace
Expand Down Expand Up @@ -107,7 +117,10 @@ DLL_PUBLIC Tensor linearize_cache_indices_cuda(
indices.scalar_type(), "linearize_cache_indices_kernel_2", [&] {
FBGEMM_LAUNCH_KERNEL(
(linearize_cache_indices_kernel<index_t, offset_t>),
div_round_up(num_indices, kMaxThreads),
utils::cuda::cap_grid_dim_x_from_workload(
num_indices,
kMaxThreads,
at::cuda::getCurrentCUDAStream()),
kMaxThreads,
0,
at::cuda::getCurrentCUDAStream(),
Expand All @@ -134,22 +147,31 @@ __launch_bounds__(kMaxThreads) void linearize_cache_indices_from_row_idx_kernel(
update_row_indices,
pta::PackedTensorAccessor32<index_t, 1, at::RestrictPtrTraits>
linear_cache_indices) {
#ifdef USE_ROCM
for (index_t index = blockIdx.x * blockDim.x + threadIdx.x;
index < update_row_indices.size(0);
index += blockDim.x * gridDim.x) {
#else
const index_t index = blockIdx.x * blockDim.x + threadIdx.x;
if (index >= update_row_indices.size(0)) {
return;
}
const int table_index = update_table_indices[index];

const auto max_offset =
::__ldg(&cache_hash_size_cumsum[cache_hash_size_cumsum.size(0) - 1]);
const auto curr_offset = ::__ldg(&cache_hash_size_cumsum[table_index]);
if (curr_offset >= 0 && update_row_indices[index] >= 0) {
linear_cache_indices[index] = update_row_indices[index] + curr_offset;
} else {
// Either table index is wrong, or index value is negative (due to pruning):
// set it to invalid value.
linear_cache_indices[index] = max_offset;
}
#endif
const int table_index = update_table_indices[index];

const auto max_offset =
::__ldg(&cache_hash_size_cumsum[cache_hash_size_cumsum.size(0) - 1]);
const auto curr_offset = ::__ldg(&cache_hash_size_cumsum[table_index]);
if (curr_offset >= 0 && update_row_indices[index] >= 0) {
linear_cache_indices[index] = update_row_indices[index] + curr_offset;
} else {
// Either table index is wrong, or index value is negative (due to
// pruning): set it to invalid value.
linear_cache_indices[index] = max_offset;
}
#ifdef USE_ROCM
} // for index (grid-stride loop, ROCm only)
#endif
}

} // namespace
Expand Down Expand Up @@ -178,7 +200,8 @@ DLL_PUBLIC Tensor linearize_cache_indices_from_row_idx_cuda(
[&] {
FBGEMM_LAUNCH_KERNEL(
(linearize_cache_indices_from_row_idx_kernel<index_t>),
div_round_up(num_indices, kMaxThreads),
utils::cuda::cap_grid_dim_x_from_workload(
num_indices, kMaxThreads, at::cuda::getCurrentCUDAStream()),
kMaxThreads,
0,
at::cuda::getCurrentCUDAStream(),
Expand Down
95 changes: 57 additions & 38 deletions fbgemm_gpu/src/split_embeddings_cache/lxu_cache.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -272,52 +272,66 @@ __global__ __launch_bounds__(kMaxThreads) void lxu_cache_lookup_kernel(
const int32_t C = lxu_cache_state.size(0);
const int32_t N =
N_unique == nullptr ? linear_cache_indices.size(0) : *N_unique;
// On ROCm the launch caps the grid to stay within the HIP 2^32
// threads-per-launch limit, so we grid-stride to cover the full workload.
// On CUDA the grid is not capped and the loop body runs once per warp.
const auto n_stride = gridDim.x * blockDim.y * blockDim.x;
#ifdef USE_ROCM
for (auto n0 =
blockIdx.x * blockDim.y * blockDim.x + threadIdx.y * blockDim.x;
n0 < N;
n0 += n_stride) {
#else
const auto n0 =
blockIdx.x * blockDim.y * blockDim.x + threadIdx.y * blockDim.x;
if (n0 >= N) {
return;
}
#endif

int32_t cache_location = kCacheLocationMissing;
int32_t n_indices = 0;
int32_t n_hits = 0;
const auto slot = threadIdx.x;
for (int i = 0; i < blockDim.x; ++i) {
const auto n = n0 + i;
if (n >= N) {
continue;
}
const int64_t idx = linear_cache_indices[n];
if (idx == invalid_index) {
continue;
}
const int32_t cache_set = cache_slot(idx, C);
n_indices++;
const bool found =
(::__ldg((&lxu_cache_state[cache_set][0]) + slot) == idx);
// fbgemm_gpu::ballot_sync wraps __ballot (ROCm, unsynchronized) vs
// __ballot_sync (CUDA). Return type is uint64_t on ROCm and uint32_t
// on CUDA, so __ffsll handles both via implicit promotion.
const auto bitmap = fbgemm_gpu::ballot_sync(found);
if (bitmap) {
// LSB == 1 hence we need to subtract one to get lane ID.
const auto way = __ffsll(static_cast<long long>(bitmap)) - 1;
if (i == threadIdx.x) {
cache_location = cache_set * kWarpSize + way;
int32_t cache_location = kCacheLocationMissing;
int32_t n_indices = 0;
int32_t n_hits = 0;
const auto slot = threadIdx.x;
for (int i = 0; i < blockDim.x; ++i) {
const auto n = n0 + i;
if (n >= N) {
continue;
}
const int64_t idx = linear_cache_indices[n];
if (idx == invalid_index) {
continue;
}
const int32_t cache_set = cache_slot(idx, C);
n_indices++;
const bool found =
(::__ldg((&lxu_cache_state[cache_set][0]) + slot) == idx);
// fbgemm_gpu::ballot_sync wraps __ballot (ROCm, unsynchronized) vs
// __ballot_sync (CUDA). Return type is uint64_t on ROCm and uint32_t
// on CUDA, so __ffsll handles both via implicit promotion.
const auto bitmap = fbgemm_gpu::ballot_sync(found);
if (bitmap) {
// LSB == 1 hence we need to subtract one to get lane ID.
const auto way = __ffsll(static_cast<long long>(bitmap)) - 1;
if (i == threadIdx.x) {
cache_location = cache_set * kWarpSize + way;
}
n_hits++;
}
n_hits++;
}
}

const auto n = n0 + threadIdx.x;
if (n < N) {
lxu_cache_locations[n] = cache_location;
}
if (gather_cache_stats && threadIdx.x == 0 && n_indices > n_hits) {
atomicAdd(
&uvm_cache_stats[uvm_cache_stats_index::num_conflict_misses],
(n_indices - n_hits));
}
const auto n = n0 + threadIdx.x;
if (n < N) {
lxu_cache_locations[n] = cache_location;
}
if (gather_cache_stats && threadIdx.x == 0 && n_indices > n_hits) {
atomicAdd(
&uvm_cache_stats[uvm_cache_stats_index::num_conflict_misses],
(n_indices - n_hits));
}
#ifdef USE_ROCM
} // for n0 (grid-stride loop, ROCm only)
#endif
}

template <typename index_t>
Expand Down Expand Up @@ -447,7 +461,12 @@ DLL_PUBLIC Tensor lxu_cache_lookup_cuda(
}

const dim3 threads(kWarpSizeHost(), kMaxThreads / kWarpSizeHost());
const dim3 blocks(div_round_up(N, kMaxThreads));
// HIP enforces a hard limit of 2^32 total threads per launch;
// lxu_cache_lookup_kernel grid-strides over n0, so capping is
// correctness-preserving. See: https://github.com/ROCm/hip/issues/2253
const dim3 blocks(
utils::cuda::cap_grid_dim_x_from_workload(
N, kMaxThreads, at::cuda::getCurrentCUDAStream()));

AT_DISPATCH_INDEX_TYPES(
linear_cache_indices.scalar_type(), "lxu_cache_lookup_cuda", [&] {
Expand Down
Loading