[CPU] improve numa memory binding (#19666)

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

docs/platforms/cpu_server.md

@@ -192,13 +192,31 @@ Notes:
     the system will automatically utilize the first `n` SNCs.
     Note that `n` cannot exceed the total SNC number, doing so will result in an error.
 
-    To specify the cores to be used, we need to explicitly set the environment variable `SGLANG_CPU_OMP_THREADS_BIND`.
+    `SGLANG_CPU_OMP_THREADS_BIND` allows explicit control of CPU cores for each tensor parallel (TP) rank.
-    For example, if we want to run the SGLang service using the first 40 cores of each SNC on a Xeon® 6980P server,
+
+    **example 1**: Run SGLang service with TP=6, using the first 40 cores of each SNC on a Xeon® 6980P server,
     which has 43-43-42 cores on the 3 SNCs of a socket, we should set:
 
     ```bash
     export SGLANG_CPU_OMP_THREADS_BIND="0-39|43-82|86-125|128-167|171-210|214-253"
     ```
+    This configuration is equivalent to:
+    - rank 0: `numactl -C 0-39 -m 0`
+    - rank 1: `numactl -C 43-82 -m 1`
+    - rank 2: `numactl -C 86-125 -m 2`
+    - rank 3: `numactl -C 128-167 -m 3`
+    - rank 4: `numactl -C 171-210 -m 4`
+    - rank 5: `numactl -C 214-253 -m 5`
+
+
+    **example 2**: Run SGLang service with TP=2, using 96 cores cross 3 SNCs on a Xeon® 6972P server,
+    which has 32-32-32 cores on the 3 SNCs in a socket, we should set:
+    ```bash
+    export SGLANG_CPU_OMP_THREADS_BIND="0-95|96-191"
+    ```
+    This configuration is equivalent to:
+    - rank 0: `numactl -C 0-95 -m 0-2`
+    - rank 1: `numactl -C 96-191 -m 3-5`
 
     Please beware that with SGLANG_CPU_OMP_THREADS_BIND set,
     the available memory amounts of the ranks may not be determined in prior.

sgl-kernel/csrc/cpu/numa_utils.cpp

@@ -30,8 +30,15 @@ std::string init_cpu_threads_env(const std::string& cpu_ids) {
 
   // Memory node binding
   if (numa_available() != -1) {
-    int mem_node_id = numa_node_of_cpu(omp_cpu_ids.front());
+    TORCH_CHECK(!omp_cpu_ids.empty(), "Cannot bind memory, no CPUs specified.");
-    bitmask* mask = numa_parse_nodestring(std::to_string(mem_node_id).c_str());
+    int mem_node_id_st = numa_node_of_cpu(omp_cpu_ids.front());
+    int mem_node_id_ed = numa_node_of_cpu(omp_cpu_ids.back());
+    if (mem_node_id_st > mem_node_id_ed) {
+      std::swap(mem_node_id_st, mem_node_id_ed);
+    }
+
+    bitmask* mask =
+        numa_parse_nodestring((std::to_string(mem_node_id_st) + "-" + std::to_string(mem_node_id_ed)).c_str());
     bitmask* src_mask = numa_get_membind();
 
     int pid = getpid();