[0xDexe]

Summary

Fixes #1570.

The experimental embedding kernels compute row offsets in int32. The forward pass multiplies offsets_m * embedding_dim and indices * embedding_dim, and the backward pass does the same. Once the output or the embedding table holds 2^31 or more elements, those products wrap around. The kernel then reads and writes out of bounds and crashes with an illegal memory access.

The fix promotes pid_m to int64 before computing start_m. It also casts the loaded indices to int64. That keeps every offset product in int64, and it works even when the caller passes int32 indices. The change is the same in embedding_forward_kernel and embedding_backward_kernel.

Details

I added two tests. Each one crosses the int32 boundary on a different side.

• test_embedding_output_offsets_beyond_int32 makes num_tokens * embedding_dim larger than 2^31. It checks the output rows past the boundary and also checks the gradient counts.
• test_embedding_table_offsets_beyond_int32_with_int32_indices makes the table larger than 2^31 elements and reads the last rows with int32 indices. Then it checks the output and the gradient.

The existing test_embedding_correctness is skip-marked as under experimentation. So the new tests call LigerEmbeddingFunction directly. They skip when the GPU has less than 24 GiB free, so they won't break smaller CI runners.

On GB10 (DGX Spark) the CPU and GPU share memory. torch.cuda.mem_get_info() reports only about 1 GiB free there because it doesn't count reclaimable page cache. So these tests skip on that machine by default. I disabled the skip check locally to run them.

Testing Done

Without the fix, both new tests fail with CUDA error: an illegal memory access was encountered. With the fix, both pass.

pytest test/transformers/test_embedding.py -k beyond_int32 -v
2 passed, 11 deselected in 2.76s

• Hardware Type: NVIDIA GB10 (DGX Spark), CUDA 13, torch 2.14.1+cu130, triton 3.8.0
• run make test to ensure correctness (ran the embedding tests above)
• run make checkstyle to ensure code style
• run make test-convergence to ensure convergence (N/A, the experimental embedding isn't used in the convergence tests)

[0xDexe]

Hi @Tcc0403, could you take a look when you get a chance? This fixes #1570. The experimental embedding kernels computed row offsets in int32, so they hit an illegal memory access once the output or the table went past 2^31 elements.
I tested it on a GB10 (DGX Spark). Both new tests crash without the fix and pass with it. They skip when the GPU has less than 24 GiB free, so they won't affect smaller CI runners.

@xiaoqiangheye this should cover the case you reported, if you want to try it.