A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 10, 22:45 UTC
[PERF] Add interleaved memory allocation to NUMA module (#30800)
skaraban3807 committed
Dec 24, 2025 at 19:17 UTC
7cd288a4b3f45f3524134b9915c35accd668c2f0
1 file changed
+38
-13
csrc/cpu/utils.cpp
+38
-13
@@ -24,6 +24,8 @@ std::string init_cpu_threads_env(const std::string& cpu_ids) {
24
#ifndef VLLM_NUMA_DISABLED
25
std::string init_cpu_threads_env(const std::string& cpu_ids) {
26
bitmask* omp_cpu_mask = numa_parse_cpustring_all(cpu_ids.c_str());
27
+ TORCH_CHECK(omp_cpu_mask != nullptr,
28
+ "Failed to parse CPU string: " + cpu_ids);
29
TORCH_CHECK(omp_cpu_mask->size > 0);
30
std::vector<int> omp_cpu_ids;
31
omp_cpu_ids.reserve(omp_cpu_mask->size);
@@ -44,20 +46,12 @@ std::string init_cpu_threads_env(const std::string& cpu_ids) {
46
47
// Memory node binding
48
if (numa_available() != -1) {
47
- int mem_node_id = numa_node_of_cpu(omp_cpu_ids.front());
49
std::set<int> node_ids;
50
for (const auto& cpu_id : omp_cpu_ids) {
51
int node_id = numa_node_of_cpu(cpu_id);
52
if (node_id != -1) {
53
node_ids.insert(node_id);
54
}
54
- if (node_id != mem_node_id) {
55
- TORCH_WARN("CPU ", cpu_id, " is on NUMA node ", node_id, ", but CPU ",
56
- omp_cpu_ids.front(), " is on NUMA node ", mem_node_id,
57
- ". All CPUs should be on the same NUMA node for optimal "
58
- "performance. Memory will be bound to NUMA node ",
59
- mem_node_id, ".");
60
- }
55
}
56
// Concatenate all node_ids into a single comma-separated string
57
if (!node_ids.empty()) {
@@ -70,7 +64,7 @@ std::string init_cpu_threads_env(const std::string& cpu_ids) {
64
}
65
66
bitmask* mask = numa_parse_nodestring(node_ids_str.c_str());
73
- bitmask* src_mask = numa_get_membind();
67
+ bitmask* src_mask = numa_get_mems_allowed();
68
69
int pid = getpid();
70
@@ -83,15 +77,46 @@ std::string init_cpu_threads_env(const std::string& cpu_ids) {
77
std::to_string(errno));
78
}
79
86
- // restrict memory allocation node.
87
- numa_set_membind(mask);
80
+ // Restrict memory allocation to the selected NUMA node(s).
81
+ // Enhances memory locality for the threads bound to those NUMA CPUs.
82
+ if (node_ids.size() > 1) {
83
+ errno = 0;
84
+ numa_set_interleave_mask(mask);
85
+ if (errno != 0) {
86
+ TORCH_WARN("numa_set_interleave_mask failed. errno: " +
87
+ std::to_string(errno));
88
+ } else {
89
+ TORCH_WARN(
90
+ "NUMA binding: Using INTERLEAVE policy for memory "
91
+ "allocation across multiple NUMA nodes (nodes: " +
92
+ node_ids_str +
93
+ "). Memory allocations will be "
94
+ "interleaved across the specified NUMA nodes.");
95
+ }
96
+ } else {
97
+ errno = 0;
98
+ numa_set_membind(mask);
99
+ if (errno != 0) {
100
+ TORCH_WARN("numa_set_membind failed. errno: " +
101
+ std::to_string(errno));
102
+ } else {
103
+ TORCH_WARN(
104
+ "NUMA binding: Using MEMBIND policy for memory "
105
+ "allocation on the NUMA nodes (" +
106
+ node_ids_str +
107
+ "). Memory allocations will be "
108
+ "strictly bound to these NUMA nodes.");
109
+ }
110
+ }
111
+
112
numa_set_strict(1);
113
114
numa_free_nodemask(mask);
115
numa_free_nodemask(src_mask);
116
} else {
93
- TORCH_WARN("numa_parse_nodestring or numa_get_membind failed. errno: " +
94
- std::to_string(errno));
117
+ TORCH_WARN(
118
+ "numa_parse_nodestring or numa_get_run_node_mask failed. errno: " +
119
+ std::to_string(errno));
120
}
121
}
122
}