From 4e23e82029a6184010045542675ebd690e4518ad Mon Sep 17 00:00:00 2001 From: Iuliu Rus Date: Tue, 21 Jul 2026 14:09:58 -0700 Subject: [PATCH] Bounds-assert table_idx and row byte range in embedding_inplace_update kernels (#6035) Summary: X-link: https://github.com/facebookresearch/FBGEMM/pull/2938 Defense-in-depth backstop for the embedding TBE in-place update scatter. `embedding_inplace_update_kernel_impl` (CUDA) and `embedding_inplace_update_cpu_kernel` index per-table metadata by `table_idx = update_table_idx[i]` and then scatter to `dev_weights[weight_offset + D_bytes * row_idx]`, neither of which was bounds-checked. An out-of-range `table_idx` reads `D_offsets` / `weights_tys` / `weights_offsets` out of bounds and then writes to a wild address; an out-of-range `row_idx` writes past the weight allocation (async aperture violation on GPU). This adds two cheap asserts, reusing existing kernel inputs (no new kernel argument): - `table_idx` is bounds-checked against `num_tables = weights_offsets.size(0)` (one entry per table; `D_offsets` has `num_tables + 1` entries) immediately after it is read, before it indexes any metadata array. - The row's byte range `[byte_offset, row_end)` is bounds-checked against the destination allocation (`dev_weights.size(0)` / `uvm_weights.size(0)`), branched by placement. The row bound is deliberately the ALLOCATION size, not a per-table `weights_offsets[table_idx + 1]`. `weights_offsets` is per-placement-buffer: under mixed DEVICE/UVM placement (the `kernel_1` path) the next table's offset can belong to the OTHER buffer, so `weights_offsets[table_idx + 1]` is not a valid end for this table and would false-abort valid callers. The exact per-table `[0, num_embeddings)` range is enforced host-side in the sigrid predictor path (parent diff D112649885); this diff is the low-level allocation backstop. Severity note: on ROCm `CUDA_KERNEL_ASSERT` compiles to a bare `abort()` (torch/headeronly/macros/Macros.h), so on AMD this is a FAIL-FAST CRASH backstop -- it converts silent cross-table / out-of-allocation corruption into a deterministic crash; it is NOT a graceful reject. The graceful reject/rollback (keeps serving on last-good weights) is the host-side check in the parent diff D112649885. The CPU kernel uses `TORCH_CHECK`, which throws a catchable exception. Blast radius: `embedding_inplace_update.cu` and `embedding_inplace_update_cpu.cpp` are shared across all fbgemm TBE inplace-update users. The asserts only fire on a genuinely out-of-range `table_idx` / `row_idx`; correct callers are unaffected, and no new kernel argument is added. Differential Revision: D112662764 --- .../embedding_inplace_update.cu | 35 +++- .../embedding_inplace_update_cpu.cpp | 54 +++-- .../embedding_inplace_update_test.cpp | 190 ++++++++++++++++-- 3 files changed, 245 insertions(+), 34 deletions(-) diff --git a/fbgemm_gpu/src/embedding_inplace_ops/embedding_inplace_update.cu b/fbgemm_gpu/src/embedding_inplace_ops/embedding_inplace_update.cu index 2f9f769efc..24a09d1b5b 100644 --- a/fbgemm_gpu/src/embedding_inplace_ops/embedding_inplace_update.cu +++ b/fbgemm_gpu/src/embedding_inplace_ops/embedding_inplace_update.cu @@ -51,6 +51,15 @@ inline __device__ void embedding_inplace_update_kernel_impl( lxu_cache_locations, const int64_t i) { const int32_t table_idx = update_table_idx[i]; + // Defense-in-depth backstop (callers are expected to validate up front). An + // out-of-range table_idx would index D_offsets / weights_tys / + // weights_offsets out of bounds and then scatter to a wild address. + // weights_offsets has one entry per table; D_offsets has num_tables + 1 + // entries, so table_idx + 1 is in range once table_idx is. + const int64_t num_tables = weights_offsets.size(0); + CUDA_KERNEL_ASSERT( + table_idx >= 0 && table_idx < num_tables && + "embedding_inplace_update: table_idx out of range"); const auto row_idx = update_row_idx[i]; // TODO: We don't need to compute these here. @@ -62,17 +71,27 @@ inline __device__ void embedding_inplace_update_kernel_impl( nbit::padded_row_size_in_bytes(D, weight_ty, row_alignment); const int64_t weight_offset = weights_offsets[table_idx]; + const int64_t byte_offset = weight_offset + + static_cast(D_bytes) * static_cast(row_idx); + const int64_t row_end = weight_offset + + static_cast(D_bytes) * (static_cast(row_idx) + 1); uint8_t* __restrict__ weight_row; if (weights_placement == PlacementType::DEVICE) { - weight_row = - &dev_weights - [weight_offset + - static_cast(D_bytes) * static_cast(row_idx)]; + // Bound the row's byte range against the dev_weights allocation. A tighter + // per-table upper bound (weights_offsets[table_idx + 1]) is NOT safe here: + // weights_offsets is per-placement-buffer, so under mixed DEVICE/UVM + // placement (kernel_1) the next entry can belong to uvm_weights and is not + // a valid end for this table. The exact per-table [0, num_embeddings) range + // is enforced host-side before dispatch; this is the allocation backstop. + CUDA_KERNEL_ASSERT( + row_idx >= 0 && byte_offset >= 0 && row_end <= dev_weights.size(0) && + "embedding_inplace_update: row scatters past dev_weights"); + weight_row = &dev_weights[byte_offset]; } else { - weight_row = - &uvm_weights - [weight_offset + - static_cast(D_bytes) * static_cast(row_idx)]; + CUDA_KERNEL_ASSERT( + row_idx >= 0 && byte_offset >= 0 && row_end <= uvm_weights.size(0) && + "embedding_inplace_update: row scatters past uvm_weights"); + weight_row = &uvm_weights[byte_offset]; } // padded_row_size_in_bytes pad each row with row_alignment (16 bytes on GPUs) diff --git a/fbgemm_gpu/src/embedding_inplace_ops/embedding_inplace_update_cpu.cpp b/fbgemm_gpu/src/embedding_inplace_ops/embedding_inplace_update_cpu.cpp index 6352cccc17..1f8ec95d6c 100644 --- a/fbgemm_gpu/src/embedding_inplace_ops/embedding_inplace_update_cpu.cpp +++ b/fbgemm_gpu/src/embedding_inplace_ops/embedding_inplace_update_cpu.cpp @@ -6,10 +6,6 @@ * LICENSE file in the root directory of this source tree. */ -#include -#include -#include - #include #include #include @@ -35,8 +31,20 @@ void embedding_inplace_update_cpu_kernel( const at::TensorAccessor& update_row_idx, const at::TensorAccessor& update_offsets, int64_t row_alignment) { + const int64_t num_tables = weights_offsets.size(0); for (const auto n : c10::irange(update_row_idx.size(0))) { int32_t t = update_table_idx[n]; + // Defense-in-depth backstop. An out-of-range table_idx would index the + // per-table metadata (D_offsets / weights_tys / weights_offsets) out of + // bounds and then scatter to a wild address. weights_offsets has one entry + // per table; D_offsets has num_tables + 1 entries. + TORCH_CHECK( + t >= 0 && t < num_tables, + "embedding_inplace_update: table index ", + t, + " out of range [0, ", + num_tables, + ")"); auto row_idx = update_row_idx[n]; SparseType weight_ty = static_cast(weights_tys[t]); @@ -46,19 +54,41 @@ void embedding_inplace_update_cpu_kernel( const int32_t D_bytes = nbit::padded_row_size_in_bytes(D, weight_ty, row_alignment); int64_t weight_offset = weights_offsets[t]; + const int64_t byte_offset = weight_offset + + static_cast(D_bytes) * static_cast(row_idx); + const int64_t row_end = weight_offset + + static_cast(D_bytes) * (static_cast(row_idx) + 1); uint8_t* __restrict__ weight_row = nullptr; const auto placement = static_cast(weights_placements[t]); + // Bound the row's byte range against the destination allocation. A tighter + // per-table bound (weights_offsets[t + 1]) is not safe: weights_offsets is + // per-placement-buffer, so under mixed HOST/UVM placement the next entry + // can belong to the other buffer. The exact per-table [0, num_embeddings) + // range is enforced by callers before dispatch; this is the allocation + // backstop. if (placement == PlacementType::HOST) { - weight_row = - &dev_weights - [weight_offset + - static_cast(D_bytes) * static_cast(row_idx)]; + TORCH_CHECK( + row_idx >= 0 && byte_offset >= 0 && row_end <= dev_weights.size(0), + "embedding_inplace_update: row_idx ", + row_idx, + " (table ", + t, + ") scatters past dev_weights of ", + dev_weights.size(0), + " bytes"); + weight_row = &dev_weights[byte_offset]; } else { - weight_row = - &uvm_weights - [weight_offset + - static_cast(D_bytes) * static_cast(row_idx)]; + TORCH_CHECK( + row_idx >= 0 && byte_offset >= 0 && row_end <= uvm_weights.size(0), + "embedding_inplace_update: row_idx ", + row_idx, + " (table ", + t, + ") scatters past uvm_weights of ", + uvm_weights.size(0), + " bytes"); + weight_row = &uvm_weights[byte_offset]; } int64_t update_weight_offset = update_offsets[n]; diff --git a/fbgemm_gpu/src/embedding_inplace_ops/embedding_inplace_update_test.cpp b/fbgemm_gpu/src/embedding_inplace_ops/embedding_inplace_update_test.cpp index 1776afca22..132a3b49c3 100644 --- a/fbgemm_gpu/src/embedding_inplace_ops/embedding_inplace_update_test.cpp +++ b/fbgemm_gpu/src/embedding_inplace_ops/embedding_inplace_update_test.cpp @@ -6,6 +6,7 @@ * LICENSE file in the root directory of this source tree. */ +#include #include #include #include @@ -28,7 +29,7 @@ int32_t get_D_bytes( } template -void test_embedding_inplace_update() { +void test_embedding_inplace_update(bool use_cpu) { int T = folly::Random::rand32() % 3 + 2; // total number of tables std::vector weight_ty_list = { SparseType::FP32, @@ -45,18 +46,34 @@ void test_embedding_inplace_update() { std::vector update_row_idx; int64_t dev_weights_offset = 0; int64_t uvm_weights_offset = 0; + // dev_weights vs uvm_weights routing differs by kernel: the CUDA kernel + // treats PlacementType::DEVICE as dev_weights, while the CPU kernel treats + // PlacementType::HOST as dev_weights. Pick the matching enum per device so + // the harness drives both buffers correctly on either path (the old + // hard-coded 0/1 convention only matched the CUDA kernel). + const int32_t devPlacement = use_cpu + ? static_cast(PlacementType::HOST) + : static_cast(PlacementType::DEVICE); + const int32_t uvmPlacement = static_cast(PlacementType::MANAGED); + // Row alignment differs by kernel (16B on device, 1B on CPU). D_bytes below + // MUST use the same alignment the kernel will use, otherwise the harness lays + // out buffers/offsets differently than the kernel scatters (the old + // hard-coded 16 only matched the device path). + const int64_t row_alignment = use_cpu ? 1L : 16L; for (const auto i : c10::irange(T)) { SparseType weight_ty = weight_ty_list[folly::Random::rand32() % weight_ty_list.size()]; weights_tys.push_back(uint8_t(weight_ty)); int D = 1 << (1 + folly::Random::rand32() % 5); // table dimension - int32_t D_bytes = nbit::padded_row_size_in_bytes(D, weight_ty, 16); + int32_t D_bytes = + nbit::padded_row_size_in_bytes(D, weight_ty, row_alignment); int total_rows = 10 + folly::Random::rand32() % 50; D_offsets.push_back(D_offsets.back() + D); - int32_t weights_placement = folly::Random::rand32() % 2; + int32_t weights_placement = + (folly::Random::rand32() % 2 == 0) ? devPlacement : uvmPlacement; weights_placements.push_back(weights_placement); - if (weights_placement == 0) { + if (weights_placement == devPlacement) { weights_offsets.push_back(dev_weights_offset); dev_weights_offset += D_bytes * total_rows; } else { @@ -83,9 +100,7 @@ void test_embedding_inplace_update() { << ", update rows: " << update_row_idx_str; } - bool use_cpu = folly::Random::rand32() % 2; auto device = use_cpu ? at::kCPU : at::kCUDA; - int64_t row_alignment = use_cpu ? 1L : 16L; auto dev_weight = at::randint( 0, 255, {dev_weights_offset}, at::device(device).dtype(at::kByte)); @@ -157,11 +172,11 @@ void test_embedding_inplace_update() { auto weight_placement = weights_placements[table_idx]; auto D = D_offsets[table_idx + 1] - D_offsets[table_idx]; SparseType ty = static_cast(weights_tys[table_idx]); - int32_t D_bytes = nbit::padded_row_size_in_bytes(D, ty, 16); + int32_t D_bytes = nbit::padded_row_size_in_bytes(D, ty, row_alignment); auto dev_weight_acc = dev_weight_cpu.const_data_ptr(); auto uvm_weight_acc = uvm_weight_cpu.const_data_ptr(); auto update_weight_acc = update_weight_cpu.const_data_ptr(); - if (weight_placement == 0) { + if (weight_placement == devPlacement) { for (const auto j : c10::irange(D_bytes)) { ASSERT_EQ( dev_weight_acc[weight_offset + D_bytes * row_idx + j], @@ -178,10 +193,157 @@ void test_embedding_inplace_update() { } } -TEST(EmbeddingInplaceUpdateTest, random_update) { - // TODO: Skipping tests for now because it is - // unreliable and crashes occasionally. This should be fixed and re-enabled. - // - // test_embedding_inplace_update(); - // test_embedding_inplace_update(); +// Positive coverage via the shared harness: a valid (in-range) random update +// must scatter correctly and must NOT trip the new bounds asserts. The harness +// randomly assigns per-table DEVICE/UVM placement, so this also exercises the +// mixed-placement path that the allocation-size row bound must not false-abort. +TEST(EmbeddingInplaceUpdateTest, CpuRandomValidUpdate) { + test_embedding_inplace_update(/*use_cpu=*/true); + test_embedding_inplace_update(/*use_cpu=*/true); +} + +// Real device-kernel (.cu) coverage: proves the CUDA_KERNEL_ASSERT guards do +// not false-abort valid callers. Skipped when no GPU is present. +// +// There is intentionally NO device NEGATIVE (assert-fires) test: on ROCm +// CUDA_KERNEL_ASSERT compiles to a bare abort(), which SIGABRTs the whole +// process and poisons the GPU/HIP context; gtest death tests fork and CUDA/HIP +// is not fork-safe, so an out-of-range device index cannot be reliably caught +// in-process (it would need the DSA framework / TORCH_USE_CUDA_DSA). The +// out-of-range negative path is instead covered at the CPU-kernel layer +// (Cpu*OutOfRange* below, TORCH_CHECK -> catchable) and the host layer +// (parent diff D112649885). +TEST(EmbeddingInplaceUpdateTest, GpuRandomValidUpdate) { + if (!at::hasCUDA()) { + GTEST_SKIP() << "No CUDA/HIP device available"; + } + test_embedding_inplace_update(/*use_cpu=*/false); + test_embedding_inplace_update(/*use_cpu=*/false); +} + +// Tests for the defense-in-depth bounds asserts added to the inplace-update +// kernels. Only the CPU kernel is exercised here: it uses TORCH_CHECK, which +// throws a catchable c10::Error. The CUDA kernel uses CUDA_KERNEL_ASSERT, which +// on ROCm compiles to a bare abort() that crashes the whole process, so the +// device out-of-range case cannot be a catchable/death gtest in-process; it is +// compile-verified only. + +namespace { + +// Apply one delta update to (tableIdx, rowIdx) of a single HOST-placement FP32 +// table with `totalRows` rows of dimension `D`, via the CPU kernel. Throws (via +// TORCH_CHECK) if a bounds assert fires. +void runSingleTableCpuUpdate( + int64_t totalRows, + int32_t D, + int32_t tableIdx, + int64_t rowIdx) { + const SparseType weight_ty = SparseType::FP32; + const int64_t row_alignment = 1; + const int32_t D_bytes = + nbit::padded_row_size_in_bytes(D, weight_ty, row_alignment); + + auto dev_weights = + at::zeros({totalRows * D_bytes}, at::device(at::kCPU).dtype(at::kByte)); + auto uvm_weights = at::zeros({0}, at::device(at::kCPU).dtype(at::kByte)); + // CPU kernel routes PlacementType::HOST to dev_weights. + auto weights_placements = at::tensor( + {static_cast(PlacementType::HOST)}, at::dtype(at::kInt)); + auto weights_offsets = at::tensor({int64_t(0)}, at::dtype(at::kLong)); + auto weights_tys = + at::tensor({static_cast(weight_ty)}, at::dtype(at::kByte)); + auto D_offsets = at::tensor({0, D}, at::dtype(at::kInt)); + auto update_weights = + at::zeros({D_bytes}, at::device(at::kCPU).dtype(at::kByte)); + auto update_table_idx = at::tensor({tableIdx}, at::dtype(at::kInt)); + auto update_row_idx = at::tensor({rowIdx}, at::dtype(at::kLong)); + auto update_offsets = + at::tensor({int64_t(0), int64_t(D_bytes)}, at::dtype(at::kLong)); + + embedding_inplace_update_cpu( + dev_weights, + uvm_weights, + weights_placements, + weights_offsets, + weights_tys, + D_offsets, + update_weights, + update_table_idx, + update_row_idx, + update_offsets, + row_alignment); +} + +} // namespace + +TEST(EmbeddingInplaceUpdateTest, CpuRejectsOutOfRangeTableIdx) { + // Single table (num_tables == 1); table index 5 is out of range and would + // index per-table metadata out of bounds. Must throw, not read wild memory. + EXPECT_ANY_THROW(runSingleTableCpuUpdate( + /*totalRows=*/8, /*D=*/4, /*tableIdx=*/5, /*rowIdx=*/0)); +} + +TEST(EmbeddingInplaceUpdateTest, CpuRejectsOutOfRangeRowIdx) { + // Row 100 in an 8-row table scatters past dev_weights. Must throw. + EXPECT_ANY_THROW(runSingleTableCpuUpdate( + /*totalRows=*/8, /*D=*/4, /*tableIdx=*/0, /*rowIdx=*/100)); +} + +TEST(EmbeddingInplaceUpdateTest, CpuValidUpdateSucceeds) { + // In-range update must not false-abort. + EXPECT_NO_THROW(runSingleTableCpuUpdate( + /*totalRows=*/8, /*D=*/4, /*tableIdx=*/0, /*rowIdx=*/5)); +} + +TEST(EmbeddingInplaceUpdateTest, CpuMixedPlacementValidUpdateSucceeds) { + // Two tables under MIXED placement: table 0 HOST (dev_weights), table 1 + // non-HOST (uvm_weights). weights_offsets is per-placement-buffer, so both + // start at offset 0 in their own buffer. A valid in-range update to each + // table must NOT abort. This locks in the allocation-size row bound: a + // per-table weights_offsets[table_idx + 1] bound would treat table 1's + // offset (0, in the uvm buffer) as table 0's end and false-abort this + // valid caller. + const SparseType weight_ty = SparseType::FP32; + const int64_t row_alignment = 1; + const int32_t D = 4; + const int32_t D_bytes = + nbit::padded_row_size_in_bytes(D, weight_ty, row_alignment); + const int64_t totalRows = 8; + + auto dev_weights = + at::zeros({totalRows * D_bytes}, at::device(at::kCPU).dtype(at::kByte)); + auto uvm_weights = + at::zeros({totalRows * D_bytes}, at::device(at::kCPU).dtype(at::kByte)); + // table 0 -> dev_weights (HOST), table 1 -> uvm_weights (non-HOST). + auto weights_placements = at::tensor( + {static_cast(PlacementType::HOST), + static_cast(PlacementType::DEVICE)}, + at::dtype(at::kInt)); + auto weights_offsets = + at::tensor({int64_t(0), int64_t(0)}, at::dtype(at::kLong)); + auto weights_tys = at::tensor( + {static_cast(weight_ty), static_cast(weight_ty)}, + at::dtype(at::kByte)); + auto D_offsets = at::tensor({0, D, 2 * D}, at::dtype(at::kInt)); + auto update_weights = + at::zeros({2 * D_bytes}, at::device(at::kCPU).dtype(at::kByte)); + auto update_table_idx = at::tensor({0, 1}, at::dtype(at::kInt)); + auto update_row_idx = + at::tensor({int64_t(5), int64_t(6)}, at::dtype(at::kLong)); + auto update_offsets = at::tensor( + {int64_t(0), int64_t(D_bytes), int64_t(2 * D_bytes)}, + at::dtype(at::kLong)); + + EXPECT_NO_THROW(embedding_inplace_update_cpu( + dev_weights, + uvm_weights, + weights_placements, + weights_offsets, + weights_tys, + D_offsets, + update_weights, + update_table_idx, + update_row_idx, + update_offsets, + row_alignment)); }