@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 39dff5ff39
Fork 0

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 3, 21:30 UTC

Add VLLM_USE_SPINLOOP_EXT to use more efficient busy polling (#36517)

Signed-off-by: Patrick Schlangen <pschlan@amd.com>

pschlan-amd committed May 12, 2026 at 01:11 UTC 39dff5ff39a8f45caeb81d345d7a3d70c2723e40
5 files changed +259 -14
CMakeLists.txt
+18
@@ -109,6 +109,24 @@ else()
109 set(CUDA_SUPPORTED_ARCHS "7.0;7.5;8.0;8.6;8.7;8.9;9.0")
110 endif()
111
112 +#
113 +# spinloop extension (pure CXX; must stay above the non-CUDA device branch so
114 +# CPU builds define the target before the early return)
115 +#
116 +set(VLLM_SPINLOOP_EXT_SRC "csrc/spinloop.cpp")
117 +set(SPINLOOP_COMPILE_FLAGS "")
118 +if(CMAKE_SYSTEM_PROCESSOR MATCHES "x86_64|amd64")
119 + list(APPEND SPINLOOP_COMPILE_FLAGS "-mmwaitx")
120 +endif()
121 +define_extension_target(
122 + spinloop
123 + DESTINATION vllm
124 + LANGUAGE CXX
125 + SOURCES ${VLLM_SPINLOOP_EXT_SRC}
126 + COMPILE_FLAGS ${SPINLOOP_COMPILE_FLAGS}
127 + USE_SABI 3.11
128 + WITH_SOABI)
129 +
130 #
131 # Forward the non-CUDA device extensions to external CMake scripts.
132 #
csrc/spinloop.cpp new
+204
@@ -0,0 +1,204 @@
1 +#include <Python.h>
2 +
3 +extern "C" {
4 +
5 +#include <stdbool.h>
6 +#include <time.h>
7 +
8 +#if defined(__i386__) || defined(__x86_64__)
9 + #include <cpuid.h>
10 + #include <mwaitxintrin.h>
11 +#endif
12 +
13 +#if defined(CLOCK_MONOTONIC_RAW)
14 + #define TIMEOUT_CLOCK CLOCK_MONOTONIC_RAW
15 +#else
16 + #define TIMEOUT_CLOCK CLOCK_MONOTONIC
17 +#endif
18 +
19 +#define CPU_SUPPORT_NONE 0
20 +#define CPU_SUPPORT_MONITORX 1
21 +
22 +#define MWAITX_DEFAULT_TIMEOUT_CYCLES 1000000
23 +
24 +typedef struct {
25 + unsigned int cpu_support;
26 + unsigned int max_monitor_line_size;
27 +} spinloop_state_t;
28 +
29 +static void determine_cpu_support(spinloop_state_t* state) {
30 + state->cpu_support = CPU_SUPPORT_NONE;
31 + state->max_monitor_line_size = 0;
32 +
33 +#if defined(__i386__) || defined(__x86_64__)
34 + unsigned int eax, ebx, ecx, edx;
35 + if (__get_cpuid(0, &eax, &ebx, &ecx, &edx) == 1) {
36 + // AMD CPU (possible monitorx/mwaitx support)
37 + if (ebx == 0x68747541 && edx == 0x69746e65 && ecx == 0x444d4163) {
38 + if (__get_cpuid(0x80000000, &eax, &ebx, &ecx, &edx) == 1 &&
39 + eax >= 0x80000001 &&
40 + __get_cpuid(0x80000001, &eax, &ebx, &ecx, &edx) == 1) {
41 + if ((ecx & (1 << 29)) != 0) {
42 + state->cpu_support = CPU_SUPPORT_MONITORX;
43 + }
44 + }
45 + }
46 + }
47 +
48 + if (state->cpu_support == CPU_SUPPORT_MONITORX) {
49 + if (__get_cpuid(5, &eax, &ebx, &ecx, &edx) == 1) {
50 + state->max_monitor_line_size = ebx & 0xff;
51 + }
52 + }
53 +#endif
54 +}
55 +
56 +static PyObject* method_spinloop(PyObject* self, PyObject* args,
57 + PyObject* kwargs) {
58 + Py_buffer buffer;
59 + PyObject* callback;
60 + double timeout = 0.;
61 +
62 + spinloop_state_t* state = (spinloop_state_t*)PyModule_GetState(self);
63 + if (state == NULL) {
64 + PyErr_SetString(PyExc_TypeError, "Failed to retrieve module state!");
65 + return NULL;
66 + }
67 +
68 + static const char* keywords[] = {"buffer", "callback", "timeout", NULL};
69 + if (!PyArg_ParseTupleAndKeywords(args, kwargs, "y*O|d", (char**)keywords,
70 + &buffer, &callback, &timeout)) {
71 + return NULL;
72 + }
73 +
74 + if (!PyCallable_Check(callback)) {
75 + PyErr_SetString(PyExc_TypeError, "callback parameter must be callable!");
76 + PyBuffer_Release(&buffer);
77 + return NULL;
78 + }
79 +
80 + struct timespec t_start;
81 + if (clock_gettime(TIMEOUT_CLOCK, &t_start) != 0) {
82 + PyErr_SetString(PyExc_RuntimeError, "clock_gettime() failed!");
83 + PyBuffer_Release(&buffer);
84 + return NULL;
85 + }
86 +
87 + bool result = false;
88 + bool error = false;
89 + bool have_timeout = (timeout > 1e-9);
90 + unsigned int iteration = 0;
91 + const bool buffer_qualifies = (buffer.len <= state->max_monitor_line_size);
92 +
93 + while (true) {
94 + PyObject* res = PyObject_CallNoArgs(callback);
95 + if (res == NULL) {
96 + error = true;
97 + break;
98 + }
99 + int ok = (res == Py_True);
100 + Py_DECREF(res);
101 +
102 + if (ok) {
103 + result = true;
104 + break;
105 + }
106 +
107 + // Check timeout at most every 16 iterations to avoid clock_gettime and
108 + // comparison cost
109 + if (have_timeout && (iteration & 15u) == 0) {
110 + struct timespec t_now;
111 + if (clock_gettime(TIMEOUT_CLOCK, &t_now) != 0) {
112 + PyErr_SetString(PyExc_RuntimeError, "clock_gettime() failed!");
113 + error = true;
114 + break;
115 + }
116 +
117 + const double elapsed = (double)(t_now.tv_sec - t_start.tv_sec) +
118 + (t_now.tv_nsec - t_start.tv_nsec) * 1e-9;
119 + if (elapsed >= timeout) {
120 + result = false;
121 + break;
122 + }
123 + }
124 + ++iteration;
125 +
126 +#if defined(__i386__) || defined(__x86_64__)
127 + // monitorx + mwaitx with qualified buffer
128 + if (buffer_qualifies && state->cpu_support == CPU_SUPPORT_MONITORX) {
129 + _mm_monitorx(buffer.buf, 0, 0);
130 +
131 + // Check once more in case the buffer has been modified while we were
132 + // arming the monitor hardware
133 + res = PyObject_CallNoArgs(callback);
134 + if (res == NULL) {
135 + error = true;
136 + break;
137 + }
138 + ok = (res == Py_True);
139 + Py_DECREF(res);
140 +
141 + if (ok) {
142 + result = true;
143 + break;
144 + }
145 +
146 + // Run mwaitx with enabled timeout (bit 1). The actual timeout value
147 + // is not very important, we just want to ensure we don't lock up
148 + // here for too long.
149 + Py_BEGIN_ALLOW_THREADS _mm_mwaitx((1 << 1), 0,
150 + MWAITX_DEFAULT_TIMEOUT_CYCLES);
151 + Py_END_ALLOW_THREADS
152 + }
153 +
154 + // Fallback: Busy poll
155 + else {
156 +#endif
157 + // Give other threads a chance to be scheduled
158 + Py_BEGIN_ALLOW_THREADS
159 +#if defined(__i386__) || defined(__x86_64__)
160 + __builtin_ia32_pause();
161 +#elif defined(__aarch64__)
162 + __asm__ volatile("yield" :: : "memory");
163 +#endif
164 + Py_END_ALLOW_THREADS
165 +#if defined(__i386__) || defined(__x86_64__)
166 + }
167 +#endif
168 + }
169 +
170 + PyBuffer_Release(&buffer);
171 +
172 + if (error) {
173 + return NULL;
174 + }
175 +
176 + if (result) {
177 + Py_RETURN_TRUE;
178 + }
179 +
180 + Py_RETURN_FALSE;
181 +}
182 +
183 +static PyMethodDef spinloop_methods[] = {
184 + {"spinloop", (PyCFunction)method_spinloop, METH_VARARGS | METH_KEYWORDS,
185 + "Wait for store with callback"},
186 + {NULL, NULL, 0, NULL}};
187 +
188 +static struct PyModuleDef spinloop_module = {
189 + PyModuleDef_HEAD_INIT, "spinloop",
190 + "Hardware-optimized spinloops for Python", sizeof(spinloop_state_t),
191 + spinloop_methods};
192 +
193 +PyMODINIT_FUNC PyInit_spinloop(void) {
194 + PyObject* m = PyModule_Create(&spinloop_module);
195 + if (m != NULL) {
196 + spinloop_state_t* state = (spinloop_state_t*)PyModule_GetState(m);
197 + if (state != NULL) {
198 + determine_cpu_support(state);
199 + }
200 + }
201 + return m;
202 +}
203 +
204 +} // extern "C"
setup.py
+3
@@ -686,6 +686,7 @@ class precompiled_wheel_utils:
686 "vllm/vllm_flash_attn/_vllm_fa2_C.abi3.so",
687 "vllm/vllm_flash_attn/_vllm_fa3_C.abi3.so",
688 "vllm/cumem_allocator.abi3.so",
689 + "vllm/spinloop.abi3.so",
690 # ROCm-specific libraries
691 "vllm/_rocm_C.abi3.so",
692 ]
@@ -993,6 +994,8 @@ if _is_cuda() or _is_hip():
994 # copying the relevant .py files from the source repository.
995 ext_modules.append(CMakeExtension(name="vllm.triton_kernels", optional=True))
996
997 +ext_modules.append(CMakeExtension(name="vllm.spinloop"))
998 +
999 if _is_hip():
1000 ext_modules.append(CMakeExtension(name="vllm._rocm_C"))
1001
vllm/distributed/device_communicators/shm_broadcast.py
+31 -14
@@ -38,6 +38,11 @@ from vllm.utils.network_utils import (
38 is_valid_ipv6_address,
39 )
40
41 +if envs.VLLM_USE_SPINLOOP_EXT:
42 + from vllm.spinloop import spinloop
43 +
44 +SPINLOOP_TIMEOUT_SECONDS = 0.1
45 +
46 if TYPE_CHECKING:
47 from _typeshed import SizedBuffer
48
@@ -540,13 +545,17 @@ class MessageQueue:
545 n_warning = 1
546 while True:
547 with self.buffer.get_metadata(self.current_idx) as metadata_buffer:
543 - # Memory fence ensures we see the latest read flags from readers.
544 - # Without this, we may read stale flags from our CPU cache and
545 - # spin indefinitely even though readers have completed.
546 - memory_fence()
547 - read_count = sum(metadata_buffer[1:])
548 - written_flag = metadata_buffer[0]
549 - if written_flag and read_count != self.buffer.n_reader:
548 +
549 + def check():
550 + memory_fence()
551 + read_count = sum(metadata_buffer[1:])
552 + written_flag = metadata_buffer[0]
553 + return not (written_flag and read_count != self.buffer.n_reader)
554 +
555 + if envs.VLLM_USE_SPINLOOP_EXT and not check():
556 + spinloop(metadata_buffer, check, timeout=SPINLOOP_TIMEOUT_SECONDS)
557 +
558 + if not check():
559 # this block is written and not read by all readers
560 # for writers, `self.current_idx` is the next block to write
561 # if this block is not ready to write,
@@ -657,13 +666,21 @@ class MessageQueue:
666 )
667 with self.buffer.get_metadata(self.current_idx) as metadata_buffer:
668 while True:
660 - # Memory fence ensures we see the latest writes from the writer.
661 - # Without this, we may read stale flags from our CPU cache
662 - # and spin indefinitely even though writer has updated them.
663 - memory_fence()
664 - read_flag = metadata_buffer[self.local_reader_rank + 1]
665 - written_flag = metadata_buffer[0]
666 - if not written_flag or read_flag:
669 +
670 + def check():
671 + memory_fence()
672 + read_flag = metadata_buffer[self.local_reader_rank + 1]
673 + written_flag = metadata_buffer[0]
674 + return not (not written_flag or read_flag)
675 +
676 + if envs.VLLM_USE_SPINLOOP_EXT and not check():
677 + spinloop(
678 + metadata_buffer[0 : self.local_reader_rank + 1],
679 + check,
680 + timeout=SPINLOOP_TIMEOUT_SECONDS,
681 + )
682 +
683 + if not check():
684 # this block is either
685 # (1) not written
686 # (2) already read by this reader
vllm/envs.py
+3
@@ -1780,6 +1780,9 @@ environment_variables: dict[str, Callable[[], Any]] = {
1780 "VLLM_LORA_ENABLE_DUAL_STREAM": lambda: bool(
1781 int(os.getenv("VLLM_LORA_ENABLE_DUAL_STREAM", "0"))
1782 ),
1783 + # If set to 1, use Python spinloop extension to poll in a more efficient
1784 + # way when using the mp backend.
1785 + "VLLM_USE_SPINLOOP_EXT": lambda: bool(int(os.getenv("VLLM_USE_SPINLOOP_EXT", "0"))),
1786 }
1787
1788