@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
cpp 149 lines 5.08 KB
Raw
1 #ifndef VLLM_NUMA_DISABLED
2 #include <numa.h>
3 #include <unistd.h>
4 #include <string>
5 #include <sched.h>
6 #endif
7 #if __GLIBC__ == 2 && __GLIBC_MINOR__ < 30
8 #include <unistd.h>
9 #include <sys/syscall.h>
10 #define gettid() syscall(SYS_gettid)
11 #endif
12
13 #include "cpu/utils.hpp"
14
15 #ifdef VLLM_NUMA_DISABLED
16 void init_cpu_memory_env(std::vector<int64_t> node_ids) {}
17 #else
18 void init_cpu_memory_env(std::vector<int64_t> node_ids) {
19 // Memory node binding
20 if (numa_available() != -1) {
21 // Concatenate all node_ids into a single comma-separated string
22 if (!node_ids.empty()) {
23 std::string node_ids_str;
24 for (const int node_id : node_ids) {
25 if (!node_ids_str.empty()) {
26 node_ids_str += ",";
27 }
28 node_ids_str += std::to_string(node_id);
29 }
30
31 bitmask* mask = numa_parse_nodestring(node_ids_str.c_str());
32 bitmask* src_mask = numa_get_mems_allowed();
33
34 int pid = getpid();
35
36 if (mask && src_mask) {
37 // move all existing pages to the specified numa node.
38 *(src_mask->maskp) = *(src_mask->maskp) ^ *(mask->maskp);
39 int page_num = numa_migrate_pages(pid, src_mask, mask);
40 if (page_num == -1) {
41 TORCH_WARN("numa_migrate_pages failed. errno: " +
42 std::to_string(errno));
43 }
44
45 // Restrict memory allocation to the selected NUMA node(s).
46 // Enhances memory locality for the threads bound to those NUMA CPUs.
47 if (node_ids.size() > 1) {
48 errno = 0;
49 numa_set_interleave_mask(mask);
50 if (errno != 0) {
51 TORCH_WARN("numa_set_interleave_mask failed. errno: " +
52 std::to_string(errno));
53 } else {
54 TORCH_WARN(
55 "NUMA binding: Using INTERLEAVE policy for memory "
56 "allocation across multiple NUMA nodes (nodes: " +
57 node_ids_str +
58 "). Memory allocations will be "
59 "interleaved across the specified NUMA nodes.");
60 }
61 } else {
62 errno = 0;
63 numa_set_membind(mask);
64 if (errno != 0) {
65 TORCH_WARN("numa_set_membind failed. errno: " +
66 std::to_string(errno));
67 } else {
68 TORCH_WARN(
69 "NUMA binding: Using MEMBIND policy for memory "
70 "allocation on the NUMA nodes (" +
71 node_ids_str +
72 "). Memory allocations will be "
73 "strictly bound to these NUMA nodes.");
74 }
75 }
76
77 numa_set_strict(1);
78
79 numa_free_nodemask(mask);
80 numa_free_nodemask(src_mask);
81 } else {
82 TORCH_WARN(
83 "numa_parse_nodestring or numa_get_run_node_mask failed. errno: " +
84 std::to_string(errno));
85 }
86 }
87 }
88 }
89 #endif // VLLM_NUMA_DISABLED
90
91 namespace cpu_utils {
92 ScratchPadManager::ScratchPadManager() : size_(0), ptr_(nullptr) {
93 this->realloc(allocation_unit * 128);
94 }
95
96 void ScratchPadManager::realloc(size_t new_size) {
97 new_size = round(new_size);
98 if (new_size > size_) {
99 void* new_ptr = std::aligned_alloc(64, new_size);
100 TORCH_CHECK(new_ptr != nullptr,
101 "ScratchPadManager: aligned_alloc failed for size ", new_size);
102 if (ptr_ != nullptr) {
103 std::free(ptr_);
104 }
105 ptr_ = new_ptr;
106 size_ = new_size;
107 }
108 }
109
110 ScratchPadManager* ScratchPadManager::get_scratchpad_manager() {
111 static ScratchPadManager manager;
112 return &manager;
113 }
114 } // namespace cpu_utils
115
116 void compute_slot_mapping_kernel_impl(const torch::Tensor query_start_loc,
117 const torch::Tensor positions,
118 const torch::Tensor block_table,
119 torch::Tensor slot_mapping,
120 const int64_t block_size) {
121 const int32_t req_num = query_start_loc.size(0) - 1;
122 const int64_t block_table_stride = block_table.stride(0);
123
124 const int32_t* __restrict__ query_start_loc_ptr =
125 query_start_loc.data_ptr<int32_t>();
126 const int64_t* __restrict__ positions_ptr = positions.data_ptr<int64_t>();
127 const int32_t* __restrict__ blocktable_ptr = block_table.data_ptr<int32_t>();
128 int64_t* __restrict__ slot_mapping_ptr = slot_mapping.data_ptr<int64_t>();
129
130 #pragma omp parallel for
131 for (int32_t req_idx = 0; req_idx < req_num; ++req_idx) {
132 int32_t token_start_idx = query_start_loc_ptr[req_idx];
133 int32_t token_end_idx = query_start_loc_ptr[req_idx + 1];
134 int32_t token_num = token_end_idx - token_start_idx;
135 const int64_t* __restrict__ curr_position_ptr =
136 positions_ptr + token_start_idx;
137 int64_t* __restrict__ curr_slot_mapping_ptr =
138 slot_mapping_ptr + token_start_idx;
139 const int32_t* __restrict__ curr_block_table_ptr =
140 blocktable_ptr + req_idx * block_table_stride;
141
142 for (int32_t token_idx = 0; token_idx < token_num; ++token_idx) {
143 int64_t token_position = curr_position_ptr[token_idx];
144 int64_t block_id = curr_block_table_ptr[token_position / block_size];
145 curr_slot_mapping_ptr[token_idx] =
146 block_id * block_size + token_position % block_size;
147 }
148 }
149 }