From 707da81e84eccc66641c0664ec8fe37ac452db8b Mon Sep 17 00:00:00 2001 From: Yuxingwang-intel Date: Mon, 7 Sep 2026 09:10:35 +0800 Subject: [PATCH] [CPU] Add native CPU kernel for MurmurHash32 (#35604) Co-authored-by: Ma Mingfei --- .../sglang/kernels/aot/csrc/cpu/sampling.cpp | 125 ++++++++++++++++++ .../aot/csrc/cpu/torch_extension_cpu.cpp | 6 + .../kernels/ops/sampling/murmur_hash.py | 9 +- test/registered/cpu/test_sampling.py | 56 ++++++++ 4 files changed, 195 insertions(+), 1 deletion(-) create mode 100644 python/sglang/kernels/aot/csrc/cpu/sampling.cpp create mode 100644 test/registered/cpu/test_sampling.py diff --git a/python/sglang/kernels/aot/csrc/cpu/sampling.cpp b/python/sglang/kernels/aot/csrc/cpu/sampling.cpp new file mode 100644 index 000000000..dce84c48b --- /dev/null +++ b/python/sglang/kernels/aot/csrc/cpu/sampling.cpp @@ -0,0 +1,125 @@ +#include +#include +#include + +#include +#include + +#include "common.h" + +namespace { +inline uint32_t rotl32(uint32_t x, int r) { + return (x << r) | (x >> (32 - r)); +} + +inline uint32_t fmix32(uint32_t h) { + h ^= h >> 16; + h *= 0x85EBCA6Bu; + h ^= h >> 13; + h *= 0xC2B2AE35u; + h ^= h >> 16; + + return h; +} + +inline uint32_t murmur3_mix(uint32_t h, uint32_t k) { + k *= 0xCC9E2D51u; + k = rotl32(k, 15); + k *= 0x1B873593u; + + h ^= k; + h = rotl32(h, 13); + h = h * 5u + 0xE6546B64u; + + return h; +} + +template +void murmur_hash32_kernel_impl( + const uint64_t* seed_ptr, + const pos_t* positions_ptr, + const int64_t* col_indices_ptr, + uint32_t* output_ptr, + int64_t n, + int64_t m) { + const int64_t total = n * m; + + at::parallel_for(0, total, 0, [&](int64_t begin, int64_t end) { + for (int64_t idx = begin; idx < end; ++idx) { + const int64_t row = idx / m; + const int64_t col_idx = idx % m; + const uint64_t seed = seed_ptr[row]; + + const uint32_t pos = static_cast(positions_ptr[row]); + + const uint32_t col = static_cast(col_indices_ptr[col_idx]); + + // Split 64-bit seed into two 32-bit blocks. + const uint32_t seed_low = static_cast(seed); + + const uint32_t seed_high = static_cast(seed >> 32); + + uint32_t h = 0; + + // Process seed_low + h = murmur3_mix(h, seed_low); + + // Process seed_high + h = murmur3_mix(h, seed_high); + + // position + h = murmur3_mix(h, pos); + + // column index + h = murmur3_mix(h, col); + + h ^= 16u; + + h = fmix32(h); + + output_ptr[idx] = h; + } + }); +} + +} // namespace + +at::Tensor murmur_hash32_cpu(const at::Tensor& seed, const at::Tensor& positions, const at::Tensor& col_indices) { + CHECK_INPUT(seed); + CHECK_INPUT(positions); + CHECK_INPUT(col_indices); + + CHECK_DIM(1, seed); + CHECK_DIM(1, positions); + CHECK_DIM(1, col_indices); + TORCH_CHECK(seed.size(0) == positions.size(0), "seed and positions must have the same length"); + + TORCH_CHECK(seed.scalar_type() == at::kUInt64, "seed must have dtype torch.uint64"); + + TORCH_CHECK( + positions.scalar_type() == at::kLong || positions.scalar_type() == at::kUInt64, + "positions must have dtype torch.int64 or torch.uint64"); + + TORCH_CHECK(col_indices.scalar_type() == at::kLong, "col_indices must have dtype torch.int64"); + + const int64_t n = seed.size(0); + const int64_t m = col_indices.size(0); + + auto output = at::empty({n, m}, seed.options().dtype(at::kUInt32)); + + const uint64_t* seed_ptr = seed.data_ptr(); + + const int64_t* col_indices_ptr = col_indices.data_ptr(); + + uint32_t* output_ptr = output.data_ptr(); + + if (n == 0 || m == 0) { + return output; + } + + AT_DISPATCH_INTEGRAL_TYPES_AND(at::ScalarType::UInt64, positions.scalar_type(), "murmur_hash32_cpu", [&] { + murmur_hash32_kernel_impl(seed_ptr, positions.data_ptr(), col_indices_ptr, output_ptr, n, m); + }); + + return output; +} diff --git a/python/sglang/kernels/aot/csrc/cpu/torch_extension_cpu.cpp b/python/sglang/kernels/aot/csrc/cpu/torch_extension_cpu.cpp index 18112310d..0100d216b 100644 --- a/python/sglang/kernels/aot/csrc/cpu/torch_extension_cpu.cpp +++ b/python/sglang/kernels/aot/csrc/cpu/torch_extension_cpu.cpp @@ -501,6 +501,8 @@ void multimodal_rotary_embedding_cpu( // CPU and memory binding std::string init_cpu_threads_env(const std::string& cpu_ids); +// murmur_hash32 +at::Tensor murmur_hash32_cpu(const at::Tensor& seed, const at::Tensor& positions, const at::Tensor& col_indices); // fused_sigmoid_gating_delta_rule_update at::Tensor fused_sigmoid_gating_delta_rule_update_cpu( const at::Tensor& A_log, @@ -887,6 +889,10 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) { // CPU and memory binding m.def("init_cpu_threads_env(str cpu_ids) -> str"); + // murmur_hash32 + m.def("murmur_hash32_cpu(Tensor seed, Tensor positions, Tensor col_indices) -> Tensor"); + m.impl("murmur_hash32_cpu", torch::kCPU, &murmur_hash32_cpu); + // fused_sigmoid_gating_delta_rule_update m.def( "fused_sigmoid_gating_delta_rule_update_cpu(Tensor A_log, Tensor dt_bias, Tensor q, Tensor k, Tensor v, Tensor " diff --git a/python/sglang/kernels/ops/sampling/murmur_hash.py b/python/sglang/kernels/ops/sampling/murmur_hash.py index 3028136b8..1c3c6ce9f 100644 --- a/python/sglang/kernels/ops/sampling/murmur_hash.py +++ b/python/sglang/kernels/ops/sampling/murmur_hash.py @@ -108,9 +108,16 @@ def murmur_hash32(seed, positions, col_indices): assert len(seed.shape) == 1 and len(col_indices.shape) == 1, ( f"Inputs must be 1D tensors {seed.shape=} {col_indices.shape=}" ) + device = seed.device + if device.type == "cpu": + import sgl_kernel # noqa: F401 + + return torch.ops.sgl_kernel.murmur_hash32_cpu.default( + seed, positions, col_indices + ) + n = seed.shape[0] m = col_indices.shape[0] - device = seed.device hashed = torch.empty((n, m), dtype=torch.uint32, device=device) BLOCK_SIZE = 1024 diff --git a/test/registered/cpu/test_sampling.py b/test/registered/cpu/test_sampling.py new file mode 100644 index 000000000..88e9236e6 --- /dev/null +++ b/test/registered/cpu/test_sampling.py @@ -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__]))