[CPU] Add native CPU kernel for MurmurHash32 (#35604)

Co-authored-by: Ma Mingfei <mingfei.ma@intel.com>
This commit is contained in:
Yuxingwang-intel
2026-09-07 09:10:35 +08:00
committed by GitHub
co-authored by Ma Mingfei
parent 30d0eb2ca9
commit 707da81e84
4 changed files with 195 additions and 1 deletions
+56
View File
@@ -0,0 +1,56 @@
import sys
import pytest
import sgl_kernel # noqa: F401
import torch
from sglang.kernels.ops.sampling.murmur_hash import murmur_hash32
from sglang.test.ci.ci_register import register_cpu_ci
register_cpu_ci(est_time=5, suite="stage-a-test-cpu-intel")
@pytest.mark.parametrize("positions_dtype", [torch.int64, torch.uint64])
def test_murmur_hash32_cpu_known_values(positions_dtype):
seed = torch.tensor([0, 1, 42, 0x123456789ABCDEF0], dtype=torch.uint64)
positions = torch.tensor([0, 7, 123, 456], dtype=positions_dtype)
col_indices = torch.tensor([0, 1, 2, 17], dtype=torch.int64)
actual = murmur_hash32(seed, positions, col_indices)
expected = torch.tensor(
[
[2167721464, 10027521, 2423355346, 2203067026],
[3755322398, 4196701286, 1002451629, 183234019],
[772287619, 548237471, 2740678348, 3656549299],
[3746406971, 2891010872, 104055988, 3372550890],
],
dtype=torch.uint32,
)
torch.testing.assert_close(actual, expected)
@pytest.mark.parametrize("positions_dtype", [torch.int64, torch.uint64])
def test_murmur_hash32_cpu_shape_and_dtype(positions_dtype):
seed = torch.tensor([1, 2, 3], dtype=torch.uint64)
positions = torch.tensor([10, 20, 30], dtype=positions_dtype)
col_indices = torch.arange(128, dtype=torch.int64)
actual = murmur_hash32(
seed,
positions,
col_indices,
)
assert actual.device.type == "cpu"
assert actual.dtype == torch.uint32
assert actual.shape == (3, 128)
if __name__ == "__main__":
sys.exit(pytest.main([__file__]))