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

Co-authored-by: Ma Mingfei <mingfei.ma@intel.com>

python/sglang/kernels/aot/csrc/cpu/sampling.cpp

@@ -0,0 +1,125 @@
+#include <ATen/ATen.h>
+#include <ATen/Parallel.h>
+#include <torch/library.h>
+
+#include <algorithm>
+#include <cstdint>
+
+#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 <typename pos_t>
+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<uint32_t>(positions_ptr[row]);
+
+      const uint32_t col = static_cast<uint32_t>(col_indices_ptr[col_idx]);
+
+      // Split 64-bit seed into two 32-bit blocks.
+      const uint32_t seed_low = static_cast<uint32_t>(seed);
+
+      const uint32_t seed_high = static_cast<uint32_t>(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<uint64_t>();
+
+  const int64_t* col_indices_ptr = col_indices.data_ptr<int64_t>();
+
+  uint32_t* output_ptr = output.data_ptr<uint32_t>();
+
+  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<scalar_t>(), col_indices_ptr, output_ptr, n, m);
+  });
+
+  return output;
+}

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 "

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

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__]))