fix(embedding): use int64 offsets in experimental embedding kernels

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
 

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