Bounds-assert table_idx and row byte range in embedding_inplace_update kernels (#6035)

[iuliur-meta]

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

[meta-codesync]

@iuliur-meta has exported this pull request. If you are a Meta employee, you can view the originating Diff in D112662764.