From bbb0b30d204e653b3598618ff2bcd11f5a0d7e5d Mon Sep 17 00:00:00 2001 From: 0xDexe Date: Tue, 10 Nov 2026 20:35:51 +0000 Subject: [PATCH] int64 offsets in experimental embedding kernels --- .../ops/experimental/embedding.py | 14 ++-- test/transformers/test_embedding.py | 65 +++++++++++++++++++ 2 files changed, 75 insertions(+), 4 deletions(-) diff --git a/src/liger_kernel/ops/experimental/embedding.py b/src/liger_kernel/ops/experimental/embedding.py index 159b9a66d..feea36a5c 100644 --- a/src/liger_kernel/ops/experimental/embedding.py +++ b/src/liger_kernel/ops/experimental/embedding.py @@ -18,11 +18,14 @@ def embedding_forward_kernel( pid_m = tl.program_id(0) pid_n = tl.program_id(1) - start_m = pid_m * BLOCK_SIZE_M + # Compute row offsets in int64: `offsets_m * embedding_dim` and + # `indices * embedding_dim` overflow int32 once the output or the + # embedding table holds 2^31 or more elements. + start_m = pid_m.to(tl.int64) * BLOCK_SIZE_M start_n = pid_n * BLOCK_SIZE_N offsets_m = start_m + tl.arange(0, BLOCK_SIZE_M) mask_m = offsets_m < n_elements - indices = tl.load(indices_ptr + offsets_m, mask=mask_m, other=0) + indices = tl.load(indices_ptr + offsets_m, mask=mask_m, other=0).to(tl.int64) offsets_n = start_n + tl.arange(0, BLOCK_SIZE_N) mask_n = offsets_n < embedding_dim @@ -50,11 +53,14 @@ def embedding_backward_kernel( pid_m = tl.program_id(0) pid_n = tl.program_id(1) - start_m = pid_m * BLOCK_SIZE_M + # Compute row offsets in int64: `offsets_m * embedding_dim` and + # `indices * embedding_dim` overflow int32 once the output or the + # embedding table holds 2^31 or more elements. + start_m = pid_m.to(tl.int64) * BLOCK_SIZE_M start_n = pid_n * BLOCK_SIZE_N offsets_m = start_m + tl.arange(0, BLOCK_SIZE_M) mask_m = offsets_m < n_elements - indices = tl.load(indices_ptr + offsets_m, mask=mask_m, other=0) + indices = tl.load(indices_ptr + offsets_m, mask=mask_m, other=0).to(tl.int64) offsets_n = start_n + tl.arange(0, BLOCK_SIZE_N) mask_n = offsets_n < embedding_dim diff --git a/test/transformers/test_embedding.py b/test/transformers/test_embedding.py index 72c3cc721..0cbd34263 100644 --- a/test/transformers/test_embedding.py +++ b/test/transformers/test_embedding.py @@ -59,3 +59,68 @@ def test_embedding_correctness(num_embeddings, embedding_dim, padding_idx, dtype liger_output.backward(grad_output) assert torch.allclose(torch_embedding.weight.grad, liger_embedding.weight.grad, atol=atol, rtol=rtol) + + +INT32_MAX = 2**31 - 1 + + +def _free_device_memory_bytes(): + if not torch.cuda.is_available(): + return 0 + free, _ = torch.cuda.mem_get_info() + return free + + +@pytest.mark.skipif( + _free_device_memory_bytes() < 24 * 1024**3, + reason="Needs about 20 GiB of free GPU memory for a >2^31-element output and its gradient", +) +def test_embedding_output_offsets_beyond_int32(): + """tokens * embedding_dim > 2^31: row offsets into output / grad_output must not wrap.""" + from liger_kernel.ops.experimental.embedding import LigerEmbeddingFunction + + num_embeddings, embedding_dim = 1024, 128 + num_tokens = INT32_MAX // embedding_dim + 1024 + assert num_tokens * embedding_dim > INT32_MAX + + torch.manual_seed(0) + weight = torch.randn(num_embeddings, embedding_dim, device=device, dtype=torch.float32, requires_grad=True) + indices = torch.randint(0, num_embeddings, (num_tokens,), device=device) + + output = LigerEmbeddingFunction.apply(weight, indices) + # Rows past the int32 boundary must hold the gathered embeddings. + tail = indices[-2048:] + torch.testing.assert_close(output[-2048:], weight.detach()[tail], rtol=0, atol=0) + + # With an all-ones gradient each weight row's gradient is its exact hit count. + output.backward(torch.ones_like(output)) + del output + counts = torch.bincount(indices, minlength=num_embeddings).to(torch.float32) + expected = counts[:, None].expand(-1, embedding_dim) + torch.testing.assert_close(weight.grad, expected, rtol=0, atol=0) + + +@pytest.mark.skipif( + _free_device_memory_bytes() < 24 * 1024**3, + reason="Needs about 17 GiB of free GPU memory for a >2^31-element table and its gradient", +) +def test_embedding_table_offsets_beyond_int32_with_int32_indices(): + """num_embeddings * embedding_dim > 2^31 with int32 indices: table offsets must not wrap.""" + from liger_kernel.ops.experimental.embedding import LigerEmbeddingFunction + + embedding_dim = 128 + num_embeddings = INT32_MAX // embedding_dim + 1024 + assert num_embeddings * embedding_dim > INT32_MAX + + weight = torch.zeros(num_embeddings, embedding_dim, device=device, dtype=torch.float32) + last_rows = torch.arange(num_embeddings - 512, num_embeddings, device=device) + weight[last_rows] = last_rows.to(torch.float32)[:, None] + weight.requires_grad_(True) + indices = last_rows.to(torch.int32) + + output = LigerEmbeddingFunction.apply(weight, indices) + torch.testing.assert_close(output, weight.detach()[last_rows], rtol=0, atol=0) + + output.backward(torch.ones_like(output)) + assert torch.all(weight.grad[last_rows] == 1) + assert int(weight.grad.sum().item()) == last_rows.numel() * embedding_dim