[Bugfix][Platform] Stop re-initializing NVML on every device-capability check (fixes #50381) (#50393)
Signed-off-by: seewoo <seewoo@ucsc.edu> Co-authored-by: Claude <noreply@anthropic.com>
Sebastian Woo committed
Aug 7, 2026 at 01:30 UTC
4f76c8ad9d8ec06e91ac9c84895e07a1913d7726
2 files changed
+82
-1
tests/cuda/test_cuda_context.py
+77
@@ -115,5 +115,82 @@ def test_get_device_capability_uses_visible_device_ordinal(monkeypatch):
115
assert seen_indices == [1]
116
117
118
+def _stub_nvml(monkeypatch) -> dict[str, int]:
119
+ """Stub NVML to report SM 9.0 and count init/shutdown pairs.
120
+
121
+ Pins `NvmlCudaPlatform` rather than the `CudaPlatform` alias: that alias is
122
+ `NonNvmlCudaPlatform` where NVML is unavailable, and that class reads
123
+ `torch.cuda` instead, which needs a real device.
124
+ """
125
+ from vllm.platforms.cuda import NvmlCudaPlatform, pynvml
126
+
127
+ calls = {"init": 0, "shutdown": 0}
128
+
129
+ monkeypatch.setattr(
130
+ pynvml, "nvmlInit", lambda: calls.__setitem__("init", calls["init"] + 1)
131
+ )
132
+ monkeypatch.setattr(
133
+ pynvml,
134
+ "nvmlShutdown",
135
+ lambda: calls.__setitem__("shutdown", calls["shutdown"] + 1),
136
+ )
137
+ # Pin the visible-device mapping so the test does not depend on whatever
138
+ # CUDA_VISIBLE_DEVICES happens to be set to in the environment.
139
+ monkeypatch.setenv(NvmlCudaPlatform.device_control_env_var, "0")
140
+ monkeypatch.setattr(
141
+ NvmlCudaPlatform,
142
+ "device_control_id_to_physical_device_id",
143
+ classmethod(lambda _cls, device_id: int(device_id)),
144
+ )
145
+ monkeypatch.setattr(
146
+ pynvml, "nvmlDeviceGetHandleByIndex", lambda index: f"handle-{index}"
147
+ )
148
+ monkeypatch.setattr(
149
+ pynvml, "nvmlDeviceGetCudaComputeCapability", lambda _handle: (9, 0)
150
+ )
151
+ NvmlCudaPlatform.get_device_capability.cache_clear()
152
+ return calls
153
+
154
+
155
+def test_has_device_capability_does_not_reinit_nvml(monkeypatch):
156
+ """Repeated capability checks must not re-enter an NVML context.
157
+
158
+ `has_device_capability` only reads the cached `get_device_capability`, which
159
+ carries its own NVML context. Wrapping it in `with_nvml_context` as well
160
+ cost an nvmlInit()/nvmlShutdown() pair per call, and
161
+ `triton_reshape_and_cache_flash` calls it per attention layer per step for
162
+ fp8 and bfloat16 KV caches (issue #50381).
163
+ """
164
+ from vllm.platforms.cuda import NvmlCudaPlatform
165
+
166
+ calls = _stub_nvml(monkeypatch)
167
+ try:
168
+ assert NvmlCudaPlatform.has_device_capability(80)
169
+ for _ in range(20):
170
+ NvmlCudaPlatform.has_device_capability(80)
171
+ NvmlCudaPlatform.has_device_capability(89)
172
+ NvmlCudaPlatform.has_device_capability((9, 0))
173
+
174
+ assert calls["init"] == 1
175
+ assert calls["shutdown"] == 1
176
+ finally:
177
+ NvmlCudaPlatform.get_device_capability.cache_clear()
178
+
179
+
180
+def test_has_device_capability_comparisons(monkeypatch):
181
+ """Dropping the redundant NVML context must not change the answers."""
182
+ from vllm.platforms.cuda import NvmlCudaPlatform
183
+
184
+ _stub_nvml(monkeypatch)
185
+ try:
186
+ assert NvmlCudaPlatform.has_device_capability(80)
187
+ assert NvmlCudaPlatform.has_device_capability(90)
188
+ assert NvmlCudaPlatform.has_device_capability((9, 0))
189
+ assert not NvmlCudaPlatform.has_device_capability(100)
190
+ assert not NvmlCudaPlatform.has_device_capability((10, 0))
191
+ finally:
192
+ NvmlCudaPlatform.get_device_capability.cache_clear()
193
+
194
+
195
if __name__ == "__main__":
196
pytest.main([__file__, "-v"])
vllm/platforms/cuda.py
+5
-1
@@ -741,12 +741,16 @@ class NvmlCudaPlatform(CudaPlatformBase):
741
return None
742
743
@classmethod
744
- @with_nvml_context
744
def has_device_capability(
745
cls,
746
capability: tuple[int, int] | int,
747
device_id: int = 0,
748
) -> bool:
749
+ # No @with_nvml_context here: the base implementation only reads
750
+ # get_device_capability(), which is cached and brings its own NVML
751
+ # context. Wrapping this method as well cost an nvmlInit()/
752
+ # nvmlShutdown() pair on every call, including calls made per attention
753
+ # layer per step from the Triton reshape-and-cache path.
754
try:
755
return super().has_device_capability(capability, device_id)
756
except RuntimeError: