| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | # ruff: noqa: E402 |
| 4 | import importlib.util |
| 5 | import os |
| 6 | |
| 7 | |
| 8 | def _get_torch_root(): |
| 9 | """Locate the installed torch package without importing it.""" |
| 10 | spec = importlib.util.find_spec("torch") |
| 11 | if not spec: |
| 12 | return None |
| 13 | if spec.origin: |
| 14 | return os.path.dirname(spec.origin) |
| 15 | if spec.submodule_search_locations: |
| 16 | return spec.submodule_search_locations[0] |
| 17 | return None |
| 18 | |
| 19 | |
| 20 | def _get_torch_version_attr(attr): |
| 21 | """Read an attribute of torch.version without importing torch. |
| 22 | |
| 23 | PyTorch version must not be determined by importing directly |
| 24 | because it will trigger the CUDA initialization, losing the |
| 25 | chance to set the LD_LIBRARY_PATH beforehand. |
| 26 | """ |
| 27 | try: |
| 28 | torch_root = _get_torch_root() |
| 29 | if not torch_root: |
| 30 | return None |
| 31 | version_path = os.path.join(torch_root, "version.py") |
| 32 | if not os.path.exists(version_path): |
| 33 | return None |
| 34 | # Load the version module without importing torch |
| 35 | ver_spec = importlib.util.spec_from_file_location("torch.version", version_path) |
| 36 | if not ver_spec or not ver_spec.loader: |
| 37 | return None |
| 38 | module = importlib.util.module_from_spec(ver_spec) |
| 39 | # Avoid registering in sys.modules to not confuse future imports |
| 40 | ver_spec.loader.exec_module(module) |
| 41 | return getattr(module, attr, None) |
| 42 | except Exception: |
| 43 | return None |
| 44 | |
| 45 | |
| 46 | def _get_torch_cuda_version(): |
| 47 | """Peripheral function to _maybe_set_cuda_compatibility_path().""" |
| 48 | return _get_torch_version_attr("cuda") |
| 49 | |
| 50 | |
| 51 | def _maybe_set_cuda_compatibility_path(): |
| 52 | """Set LD_LIBRARY_PATH for CUDA forward compatibility if enabled. |
| 53 | |
| 54 | Must run before 'import torch' since torch loads CUDA shared libraries |
| 55 | at import time and the dynamic linker only consults LD_LIBRARY_PATH when |
| 56 | a library is first loaded. |
| 57 | |
| 58 | CUDA forward compatibility is only supported on select professional and |
| 59 | datacenter NVIDIA GPUs. Consumer GPUs (GeForce, RTX) do not support it |
| 60 | and will get Error 803 if compat libs are loaded. |
| 61 | """ |
| 62 | enable = os.environ.get("VLLM_ENABLE_CUDA_COMPATIBILITY", "0").strip().lower() in ( |
| 63 | "1", |
| 64 | "true", |
| 65 | ) |
| 66 | if not enable: |
| 67 | return |
| 68 | |
| 69 | cuda_compat_path = os.environ.get("VLLM_CUDA_COMPATIBILITY_PATH", "") |
| 70 | if not cuda_compat_path or not os.path.isdir(cuda_compat_path): |
| 71 | conda_prefix = os.environ.get("CONDA_PREFIX", "") |
| 72 | conda_compat = os.path.join(conda_prefix, "cuda-compat") |
| 73 | if conda_prefix and os.path.isdir(conda_compat): |
| 74 | cuda_compat_path = conda_compat |
| 75 | if not cuda_compat_path or not os.path.isdir(cuda_compat_path): |
| 76 | torch_cuda_version = _get_torch_cuda_version() |
| 77 | if torch_cuda_version: |
| 78 | default_path = f"/usr/local/cuda-{torch_cuda_version}/compat" |
| 79 | if os.path.isdir(default_path): |
| 80 | cuda_compat_path = default_path |
| 81 | if not cuda_compat_path or not os.path.isdir(cuda_compat_path): |
| 82 | return |
| 83 | |
| 84 | norm_path = os.path.normpath(cuda_compat_path) |
| 85 | existing = os.environ.get("LD_LIBRARY_PATH", "") |
| 86 | ld_paths = existing.split(os.pathsep) if existing else [] |
| 87 | |
| 88 | if ld_paths and ld_paths[0] and os.path.normpath(ld_paths[0]) == norm_path: |
| 89 | return # Already at the front |
| 90 | |
| 91 | new_paths = [norm_path] + [ |
| 92 | p for p in ld_paths if not p or os.path.normpath(p) != norm_path |
| 93 | ] |
| 94 | os.environ["LD_LIBRARY_PATH"] = os.pathsep.join(new_paths) |
| 95 | |
| 96 | |
| 97 | def _maybe_promote_torch_symbols_for_rocm(): |
| 98 | """Put libtorch_cpu.so in the global symbol scope on ROCm, for GPU profiling. |
| 99 | |
| 100 | Must run before 'import torch'. libkineto advertises itself to |
| 101 | rocprofiler-sdk by exporting rocprofiler_configure from libtorch_cpu.so, |
| 102 | and rocprofiler-sdk looks for that symbol exactly once, from a lazy |
| 103 | initializer that runs *during* 'import torch'. CPython loads extension |
| 104 | modules RTLD_LOCAL, so unless something has already placed the symbol in |
| 105 | the global scope the lookup misses and no client is ever registered. |
| 106 | |
| 107 | The failure is silent and total: no queue interception, so roctracer |
| 108 | yields no GPU records, while torch.profiler still writes a complete |
| 109 | CPU-only trace and reports success. |
| 110 | |
| 111 | Best effort. If the library cannot be loaded early we leave the process |
| 112 | exactly as it would have been, losing only GPU tracing. |
| 113 | """ |
| 114 | if _get_torch_version_attr("hip") is None: |
| 115 | return |
| 116 | try: |
| 117 | import ctypes |
| 118 | import glob |
| 119 | |
| 120 | torch_root = _get_torch_root() |
| 121 | if not torch_root: |
| 122 | return |
| 123 | lib = os.path.join(torch_root, "lib", "libtorch_cpu.so") |
| 124 | if not os.path.exists(lib): |
| 125 | return |
| 126 | try: |
| 127 | ctypes.CDLL(lib, mode=ctypes.RTLD_GLOBAL) |
| 128 | except OSError: |
| 129 | # Some ROCm wheel layouts leave a stale RUNPATH in libtorch_cpu.so, |
| 130 | # so a dependency that resolves once torch has set things up does |
| 131 | # not resolve for a standalone load. Satisfy it by search, then |
| 132 | # retry; give up quietly if it is not where we expect. |
| 133 | site_root = os.path.dirname(torch_root) |
| 134 | deps = glob.glob( |
| 135 | os.path.join( |
| 136 | site_root, "_rocm_sdk_*", "lib", "*", "lib", "librocm-openblas.so.*" |
| 137 | ) |
| 138 | ) |
| 139 | if not deps: |
| 140 | return |
| 141 | ctypes.CDLL(deps[0], mode=ctypes.RTLD_LOCAL) |
| 142 | ctypes.CDLL(lib, mode=ctypes.RTLD_GLOBAL) |
| 143 | except Exception: |
| 144 | return |
| 145 | |
| 146 | |
| 147 | _maybe_set_cuda_compatibility_path() |
| 148 | _maybe_promote_torch_symbols_for_rocm() |
| 149 | |
| 150 | import torch |
| 151 | |
| 152 | from vllm.logger import init_logger |
| 153 | from vllm.utils.torch_utils import is_torch_equal, is_torch_equal_or_newer |
| 154 | |
| 155 | logger = init_logger(__name__) |
| 156 | |
| 157 | # set some common config/environment variables that should be set |
| 158 | # for all processes created by vllm and all processes |
| 159 | # that interact with vllm workers. |
| 160 | # they are executed whenever `import vllm` is called. |
| 161 | |
| 162 | # see https://github.com/vllm-project/vllm/pull/15951 |
| 163 | # it avoids unintentional cuda initialization from torch.cuda.is_available() |
| 164 | os.environ["PYTORCH_NVML_BASED_CUDA_CHECK"] = "1" |
| 165 | |
| 166 | # see https://github.com/vllm-project/vllm/issues/10480 and |
| 167 | # https://github.com/vllm-project/vllm/issues/10619. |
| 168 | os.environ["TORCHINDUCTOR_COMPILE_THREADS"] = "1" |
| 169 | |
| 170 | # Enable Triton autotuning result caching to disk by default. |
| 171 | # Without this, Triton re-runs autotuning on every process restart, |
| 172 | # adding significant latency to the first inference request. |
| 173 | # This writes autotuning results to TRITON_CACHE_DIR. |
| 174 | # It can still be overridden by setting TRITON_CACHE_AUTOTUNING=0 |
| 175 | # in the environment. |
| 176 | os.environ.setdefault("TRITON_CACHE_AUTOTUNING", "1") |
| 177 | |
| 178 | # When unset, TileLang routes JIT temp dirs through a world-shared |
| 179 | # /tmp/tvm-debug-mode-tempdirs/ whose ownership is pinned to whichever |
| 180 | # user compiled first, breaking every other user on a shared host. |
| 181 | # Opt into per-process tempdirs unless the user explicitly chose the |
| 182 | # debug layout (see https://github.com/vllm-project/vllm/issues/41410). |
| 183 | os.environ.setdefault("TILELANG_CLEANUP_TEMP_FILES", "1") |
| 184 | |
| 185 | # =================================================== |
| 186 | # torch 2.9 Inductor PythonWrapperCodegen monkeypatch |
| 187 | # =================================================== |
| 188 | # This change monkeypatches memory_plan_reuse in pytorch 2.9.0 to work around |
| 189 | # a test failure for test_multi_graph_piecewise_compile_outputs_equal. |
| 190 | # For more context, see https://github.com/pytorch/pytorch/pull/165514. |
| 191 | |
| 192 | |
| 193 | def memory_plan_reuse_patched(self): |
| 194 | import torch._inductor.ir as ir |
| 195 | from torch._inductor.codegen.wrapper import ( |
| 196 | EnterSubgraphLine, |
| 197 | ExitSubgraphLine, |
| 198 | MemoryPlanningLine, |
| 199 | MemoryPlanningState, |
| 200 | SubgraphPythonWrapperCodegen, |
| 201 | ) |
| 202 | from torch._inductor.virtualized import V |
| 203 | |
| 204 | def get_output_names(graph_outputs) -> list[str]: |
| 205 | import itertools |
| 206 | |
| 207 | names = [] |
| 208 | shape_counter = itertools.count(0) |
| 209 | none_counter = itertools.count(0) |
| 210 | for node in graph_outputs: |
| 211 | if isinstance(node, ir.NoneAsConstantBuffer): |
| 212 | names.append(f"{V.graph.name}_none{next(none_counter)}") |
| 213 | elif isinstance(node, ir.ShapeAsConstantBuffer): |
| 214 | names.append(f"{V.graph.name}_shape{next(shape_counter)}") |
| 215 | else: |
| 216 | names.append(node.get_name()) |
| 217 | return names |
| 218 | |
| 219 | if ( |
| 220 | isinstance(V.graph.wrapper_code, SubgraphPythonWrapperCodegen) |
| 221 | and V.graph.wrapper_code.partition_signatures is not None |
| 222 | ): |
| 223 | out_names = get_output_names( |
| 224 | V.graph.wrapper_code.partition_signatures.output_nodes |
| 225 | ) |
| 226 | else: |
| 227 | out_names = V.graph.get_output_names() |
| 228 | |
| 229 | while ( |
| 230 | self.lines |
| 231 | and isinstance(self.lines[-1], MemoryPlanningLine) |
| 232 | and self.lines[-1].node.name not in out_names # type: ignore[attr-defined] |
| 233 | ): |
| 234 | # these lines will be pointless |
| 235 | self.lines.pop() |
| 236 | |
| 237 | # codegen allocations in two passes |
| 238 | planning_states = [MemoryPlanningState()] |
| 239 | past_planning_states = [] |
| 240 | for i in range(len(self.lines)): |
| 241 | line = self.lines[i] |
| 242 | if isinstance(line, MemoryPlanningLine): |
| 243 | self.lines[i] = line.plan(planning_states[-1]) |
| 244 | elif isinstance(line, EnterSubgraphLine): |
| 245 | planning_states.append(MemoryPlanningState()) |
| 246 | elif isinstance(line, ExitSubgraphLine): |
| 247 | past_planning_states.append(planning_states.pop()) |
| 248 | past_planning_states.append(planning_states.pop()) |
| 249 | assert len(planning_states) == 0 |
| 250 | |
| 251 | |
| 252 | # =================================================== |
| 253 | # torch 2.9 Inductor get_graph_partition_signature monkeypatch |
| 254 | # =================================================== |
| 255 | # This change monkeypatches get_graph_partition_signature in pytorch 2.9.0 to |
| 256 | # fix inductor partition + attention-nvfp4 quant fusion, tested in |
| 257 | # `tests/compile/test_fusion_attn.py::test_attn_quant`. |
| 258 | # For more context, see https://github.com/pytorch/pytorch/pull/165815. |
| 259 | |
| 260 | |
| 261 | def get_graph_partition_signature_patched( |
| 262 | self, partitions, skip_cudagraphs: list[bool] |
| 263 | ): |
| 264 | """Gets signature for each graph partition, including input nodes, output nodes, and |
| 265 | whether deallocating an input within graph partition. |
| 266 | """ |
| 267 | from torch._inductor import dependencies |
| 268 | from torch._inductor.ir import GraphPartitionSignature, MutationOutput, NoneLayout |
| 269 | from torch._inductor.virtualized import V |
| 270 | from torch.utils._ordered_set import OrderedSet |
| 271 | |
| 272 | signatures = [] |
| 273 | |
| 274 | unmet_output_names = OrderedSet(V.graph.get_output_names()) |
| 275 | name_to_node = self.get_name_to_nodes() |
| 276 | |
| 277 | def is_none_layout(buf_name: str) -> bool: |
| 278 | """Checks if buf_name is NoneLayout. Buffers with NoneLayout is not allocated |
| 279 | so graph partition should not take it as inputs or outputs. |
| 280 | """ |
| 281 | buf = self.name_to_buf.get(buf_name, None) |
| 282 | |
| 283 | if buf is None: |
| 284 | return False |
| 285 | |
| 286 | if isinstance(buf.node.layout, NoneLayout): |
| 287 | if isinstance(buf.node, MutationOutput) and ( |
| 288 | real_name := self.mutation_real_name.get(buf_name, None) |
| 289 | ): |
| 290 | return is_none_layout(real_name) |
| 291 | |
| 292 | return True |
| 293 | |
| 294 | return False |
| 295 | |
| 296 | for partition, skip_cudagraph in zip( |
| 297 | reversed(partitions), reversed(skip_cudagraphs) |
| 298 | ): |
| 299 | output_names: OrderedSet[str] = OrderedSet() |
| 300 | |
| 301 | for node in partition: |
| 302 | output_names.update(node.outputs_by_name.keys()) |
| 303 | |
| 304 | returned_output_names = output_names.intersection(unmet_output_names) |
| 305 | |
| 306 | # all reads/writes are partition inputs except those generated |
| 307 | # within the partition and tensor constants |
| 308 | read_writes = dependencies.ReadWrites.merge_list( |
| 309 | [node.read_writes for node in partition] |
| 310 | ) |
| 311 | |
| 312 | # WeakDep is fake dependency on unused buffer. It should not appear |
| 313 | # in partition_input_names for inputs that are actually read or written. |
| 314 | partition_input_names = ( |
| 315 | OrderedSet( |
| 316 | [ |
| 317 | x.name |
| 318 | for x in read_writes.reads | read_writes.writes |
| 319 | if not is_none_layout(x.name) |
| 320 | ] |
| 321 | ) |
| 322 | - output_names |
| 323 | ) |
| 324 | |
| 325 | partition_input_names = OrderedSet( |
| 326 | self.mutation_real_name.get(name, name) for name in partition_input_names |
| 327 | ) |
| 328 | |
| 329 | buffer_names_to_free: OrderedSet[str] = OrderedSet() |
| 330 | for node in partition: |
| 331 | buffer_names_to_free.update(node.last_usage) |
| 332 | |
| 333 | # buffer_names_to_free may contain buffers allocated in previous |
| 334 | # graph partitions. These buffers should also be a partition |
| 335 | # input. |
| 336 | extra_input_names = [ |
| 337 | name |
| 338 | for name in (buffer_names_to_free - output_names) |
| 339 | if name in name_to_node |
| 340 | ] |
| 341 | partition_input_names.update(extra_input_names) |
| 342 | |
| 343 | input_nodes = { |
| 344 | name: name_to_node[name] |
| 345 | for name in partition_input_names |
| 346 | if name in name_to_node |
| 347 | } |
| 348 | input_deallocation = { |
| 349 | name: name in buffer_names_to_free |
| 350 | for name in partition_input_names |
| 351 | if name in name_to_node |
| 352 | } |
| 353 | |
| 354 | # if an input tensor is not freed in the partition function, it should |
| 355 | # also be returned as an output. This brings benefits to cudagraph |
| 356 | # since the returned output tensor is a cudagraph managed tensor with |
| 357 | # a static tensor address. |
| 358 | extra_output_names = [ |
| 359 | name |
| 360 | for name in partition_input_names |
| 361 | if name in name_to_node and name not in buffer_names_to_free |
| 362 | ] |
| 363 | |
| 364 | returned_output_names.update(extra_output_names) |
| 365 | |
| 366 | returned_output_names = OrderedSet( |
| 367 | self.mutation_real_name.get(name, name) for name in returned_output_names |
| 368 | ) |
| 369 | |
| 370 | output_nodes = [ |
| 371 | name_to_node[name] |
| 372 | for name in returned_output_names |
| 373 | if not is_none_layout(name) |
| 374 | ] |
| 375 | |
| 376 | constant_names = [ |
| 377 | name for name in partition_input_names if name in V.graph.constants |
| 378 | ] |
| 379 | |
| 380 | symbol_inputs = self.get_graph_partition_symbol_inputs(partition, input_nodes) |
| 381 | |
| 382 | partition_signature = GraphPartitionSignature( |
| 383 | symbol_inputs, |
| 384 | input_nodes, |
| 385 | output_nodes, |
| 386 | input_deallocation, |
| 387 | skip_cudagraph, |
| 388 | constant_names, |
| 389 | ) |
| 390 | |
| 391 | signatures.append(partition_signature) |
| 392 | |
| 393 | unmet_output_names = partition_input_names.union( |
| 394 | unmet_output_names - returned_output_names |
| 395 | ) |
| 396 | |
| 397 | return signatures[::-1] |
| 398 | |
| 399 | |
| 400 | # ======================================== |
| 401 | # torch 2.9 Inductor Scheduler monkeypatch |
| 402 | # ======================================== |
| 403 | # This change monkeypatches a function in Inductor to work around the following |
| 404 | # bug: https://github.com/vllm-project/vllm/issues/26678 |
| 405 | # |
| 406 | # The bug occurs when `use_inductor_graph_partition` is turned on and there |
| 407 | # exists operators inside of `splitting_ops` that have an in-place mutation. In |
| 408 | # vllm, this specifically occurs on the operator |
| 409 | # vllm.unified_attention_with_output. In this case, inductor does not populate |
| 410 | # the inductor IR's `origin_node` field, causing an assertion error when trying |
| 411 | # to access the node's `origin_node` field. |
| 412 | # |
| 413 | # So, we will monkeypatch torch._inductor.scheduler.Scheduler.should_partition |
| 414 | # so that it does not access the inductor IR node's `origin_node` field and just |
| 415 | # returns True if a node is registered as having a custom partition function. |
| 416 | # This is ok for now since vllm's implementation of the custom partition |
| 417 | # functions just return True. |
| 418 | # ======================================== |
| 419 | |
| 420 | |
| 421 | def should_partition_patched(self, node, should_log: bool = False) -> bool: |
| 422 | # This is a patched version of |
| 423 | # torch._inductor.scheduler.Scheduler.should_partition that modifies |
| 424 | # the following piece of code so that we always return True: |
| 425 | # https://github.com/pytorch/pytorch/blob/ecb53078faf86ca1b33277df33b82985675bb011/torch/_inductor/scheduler.py#L4712-L4724 |
| 426 | """Return True if we should partition the inductor graph on this node.""" |
| 427 | import torch._inductor.ir as ir |
| 428 | from torch._inductor.scheduler import ( |
| 429 | BaseSchedulerNode, |
| 430 | FusedSchedulerNode, |
| 431 | ) |
| 432 | from torch._inductor.utils import ( |
| 433 | _unstable_customized_partition_wrapper, |
| 434 | is_cudagraph_unsafe_op, |
| 435 | maybe_log_cudagraph_partition, |
| 436 | ) |
| 437 | |
| 438 | # Allow users to manually specify if a node should be partitioned |
| 439 | # Can only do this for FallbackKernels |
| 440 | ir_node = node.node |
| 441 | if isinstance(ir_node, torch._inductor.ir.FallbackKernel) and ( |
| 442 | op := ir_node.op_overload |
| 443 | ): |
| 444 | op_overload_packet_name = op.name() |
| 445 | op_overload_name = ( |
| 446 | f"{op_overload_packet_name}.{op._overloadname}" |
| 447 | if isinstance(op, torch._ops.OpOverload) |
| 448 | else op_overload_packet_name |
| 449 | ) |
| 450 | if ( |
| 451 | op_overload_packet_name |
| 452 | in torch._inductor.config.custom_should_partition_ops |
| 453 | or op_overload_name in torch._inductor.config.custom_should_partition_ops |
| 454 | ): |
| 455 | assert isinstance(op, torch._ops.OpOverload) |
| 456 | return True |
| 457 | |
| 458 | # When not using cudagraphs, keep all kernels in the `call` function |
| 459 | # instead of graph partition functions, since graph partition only brings |
| 460 | # benefit to cudagraph |
| 461 | if ( |
| 462 | not torch._inductor.config.triton.cudagraphs |
| 463 | and _unstable_customized_partition_wrapper.wrapper is None |
| 464 | ): |
| 465 | return True |
| 466 | |
| 467 | # avoid duplicating logs when should_partition is called multiple times |
| 468 | # on the same node |
| 469 | def noop_log(msg: str, node: BaseSchedulerNode | None) -> None: |
| 470 | return |
| 471 | |
| 472 | log_partition_reason = maybe_log_cudagraph_partition if should_log else noop_log |
| 473 | |
| 474 | if isinstance(node, FusedSchedulerNode): |
| 475 | return any(self.should_partition(snode) for snode in node.snodes) |
| 476 | |
| 477 | assert node.node is not None |
| 478 | |
| 479 | if not node.is_gpu(): |
| 480 | log_partition_reason("non gpu ops", node=node) |
| 481 | |
| 482 | return True |
| 483 | |
| 484 | if isinstance(node.node, ir.DeviceCopy): |
| 485 | log_partition_reason("DeviceCopy ops", node=node) |
| 486 | return True |
| 487 | |
| 488 | if isinstance(node.node, ir.Conditional): |
| 489 | log_partition_reason("Conditional ops", node=node) |
| 490 | return True |
| 491 | |
| 492 | if getattr(node.node, "unbacked_bindings", None): |
| 493 | log_partition_reason("unbacked binding ops", node=node) |
| 494 | return True |
| 495 | |
| 496 | if is_cudagraph_unsafe_op(node.node): |
| 497 | log_partition_reason("CUDAGraph-unsafe custom ops", node=node) |
| 498 | return True |
| 499 | |
| 500 | return False |
| 501 | |
| 502 | |
| 503 | def _update_scheduler_patched(self) -> None: |
| 504 | # Copied from torch._inductor.graph.GrahLowering._update_scheduler. Patches |
| 505 | # this method so that we can patch Scheduler.should_partition with the |
| 506 | # function above |
| 507 | """(Re)initializes the scheduler member. When initializing the scheduler, no CUBIN |
| 508 | files should be generated (to avoid biasing any benchmarks and pessimizing |
| 509 | fusion decisions). |
| 510 | """ |
| 511 | import torch._inductor.config as config |
| 512 | from torch._inductor.scheduler import Scheduler |
| 513 | |
| 514 | Scheduler.should_partition = should_partition_patched |
| 515 | Scheduler.get_graph_partition_signature = get_graph_partition_signature_patched |
| 516 | |
| 517 | with config.patch("triton.store_cubin", False): |
| 518 | self.scheduler = Scheduler(self.operations) |
| 519 | |
| 520 | |
| 521 | # =================================================== |
| 522 | # torch 2.9 Inductor get_raw_stream workaround |
| 523 | # =================================================== |
| 524 | # Workaround for TorchInductor autotune using get_raw_stream() without defining it. |
| 525 | # This occurs when compile_sizes > 1 in compilation_config. |
| 526 | # For more context, see https://github.com/vllm-project/vllm/issues/30905. |
| 527 | def _patch_get_raw_stream_if_needed(): |
| 528 | """Workaround for TorchInductor autotune get_raw_stream() bug.""" |
| 529 | from vllm.utils.torch_utils import is_torch_equal |
| 530 | |
| 531 | # Only apply the patch for torch 2.9.0 or 2.9.1 |
| 532 | if is_torch_equal("2.9.0") or is_torch_equal("2.9.1"): |
| 533 | import builtins |
| 534 | |
| 535 | # Check if CUDA functionality is available without initializing CUDA |
| 536 | # _cuda_getCurrentRawStream only exists in CUDA builds of PyTorch |
| 537 | if hasattr(torch._C, "_cuda_getCurrentRawStream"): |
| 538 | from torch._C import _cuda_getCurrentRawStream as _get_raw_stream |
| 539 | |
| 540 | builtins.get_raw_stream = _get_raw_stream # type: ignore[attr-defined] |
| 541 | |
| 542 | |
| 543 | _patch_get_raw_stream_if_needed() |
| 544 | |
| 545 | if is_torch_equal("2.9.0"): |
| 546 | from torch._inductor.codegen.wrapper import PythonWrapperCodegen |
| 547 | from torch._inductor.graph import GraphLowering |
| 548 | from torch.utils._config_module import _Config, _ConfigEntry |
| 549 | |
| 550 | # `custom_should_partition_ops` is a new config after 2.9.0. So this would |
| 551 | # not overwrite any user configs. |
| 552 | torch._inductor.config._config["custom_should_partition_ops"] = _ConfigEntry( |
| 553 | _Config(default=[]) |
| 554 | ) |
| 555 | |
| 556 | PythonWrapperCodegen.memory_plan_reuse = memory_plan_reuse_patched |
| 557 | GraphLowering._update_scheduler = _update_scheduler_patched |
| 558 | |
| 559 | # =================================================== |
| 560 | # torch <2.12 GraphCaptureOutput.get_runtime_env monkeypatch |
| 561 | # =================================================== |
| 562 | # PyTorch's AOT compile path omits builtins from used_globals, causing |
| 563 | # 'Missing required external references' errors for refs like 'type'. |
| 564 | # (which happens in transformers code) |
| 565 | # This mirrors the fix in https://github.com/pytorch/pytorch/pull/177558 |
| 566 | # and can be removed once torch >=2.12 is the minimum supported version. |
| 567 | |
| 568 | # =================================================== |
| 569 | # torch >= 2.11 Inductor constrain_to_fx_strides monkeypatch |
| 570 | # =================================================== |
| 571 | # Inductor's constrain_to_fx_strides calls .stride() on every FX arg's meta |
| 572 | # value, which crashes on FakeScriptObject (the compile-time proxy for hoisted |
| 573 | # opaque types). The patched version skips args whose meta value is not a |
| 574 | # torch.Tensor. |
| 575 | # Upstream issue: https://github.com/pytorch/pytorch/issues/175973 |
| 576 | |
| 577 | |
| 578 | _constrain_to_fx_strides_patched = False |
| 579 | |
| 580 | |
| 581 | def _apply_constrain_to_fx_strides_patch(): |
| 582 | """Patch lowering.constrain_to_fx_strides globally. Safe to call |
| 583 | multiple times; only the first call does anything. |
| 584 | Only applies for torch >= 2.11 and < 2.12.""" |
| 585 | global _constrain_to_fx_strides_patched |
| 586 | if _constrain_to_fx_strides_patched: |
| 587 | return |
| 588 | _constrain_to_fx_strides_patched = True |
| 589 | |
| 590 | if not is_torch_equal_or_newer("2.11.0.dev") or is_torch_equal_or_newer( |
| 591 | "2.12.0.dev" |
| 592 | ): |
| 593 | return |
| 594 | |
| 595 | import torch._inductor.ir as _ir |
| 596 | import torch._inductor.lowering as _lowering |
| 597 | from torch._inductor.virtualized import V as _V |
| 598 | |
| 599 | def _patched(fx_node, *args, **kwargs): |
| 600 | def apply_constraint(arg, fx_arg): |
| 601 | if isinstance(arg, _ir.IRNode): |
| 602 | meta_val = fx_arg.meta.get("val") |
| 603 | if isinstance(meta_val, torch.Tensor): |
| 604 | stride_order = _ir.get_stride_order( |
| 605 | meta_val.stride(), _V.graph.sizevars.shape_env |
| 606 | ) |
| 607 | return _ir.ExternKernel.require_stride_order(arg, stride_order) |
| 608 | return arg |
| 609 | if isinstance(arg, dict): |
| 610 | return {key: apply_constraint(arg[key], fx_arg[key]) for key in arg} |
| 611 | return arg |
| 612 | |
| 613 | args = tuple( |
| 614 | apply_constraint(arg, fx_arg) for arg, fx_arg in zip(args, fx_node.args) |
| 615 | ) |
| 616 | kwargs = {k: apply_constraint(v, fx_node.kwargs[k]) for k, v in kwargs.items()} |
| 617 | return args, kwargs |
| 618 | |
| 619 | _lowering.constrain_to_fx_strides = _patched |
| 620 | |
| 621 | |
| 622 | if is_torch_equal_or_newer("2.10.0") and not is_torch_equal_or_newer("2.12.0.dev"): |
| 623 | import builtins as _builtins |
| 624 | import pickle |
| 625 | |
| 626 | from torch._dynamo.convert_frame import GraphCaptureOutput |
| 627 | |
| 628 | _original_get_runtime_env = GraphCaptureOutput.get_runtime_env |
| 629 | |
| 630 | def _safe_builtins_dict(builtins_dict: dict) -> dict: |
| 631 | """Filter a builtins dict to only picklable entries for serialization.""" |
| 632 | result = {} |
| 633 | for k, v in builtins_dict.items(): |
| 634 | try: |
| 635 | pickle.dumps(v) |
| 636 | result[k] = v |
| 637 | except Exception: |
| 638 | pass |
| 639 | return result |
| 640 | |
| 641 | def _patched_get_runtime_env(self): # type: ignore[no-untyped-def] |
| 642 | runtime_env = _original_get_runtime_env(self) |
| 643 | for ref in runtime_env.external_refs: |
| 644 | if ref not in runtime_env.used_globals: |
| 645 | if ref.startswith("__builtins_dict__") and ref in self.f_globals: |
| 646 | runtime_env.used_globals[ref] = _safe_builtins_dict( |
| 647 | self.f_globals[ref] |
| 648 | ) |
| 649 | elif hasattr(_builtins, ref): |
| 650 | runtime_env.used_globals[ref] = getattr(_builtins, ref) |
| 651 | return runtime_env |
| 652 | |
| 653 | GraphCaptureOutput.get_runtime_env = _patched_get_runtime_env |
| 654 | |
| 655 | # =================================================== |
| 656 | # torch 2.10 FxGraphCachePickler.dumps ValueError fix |
| 657 | # =================================================== |
| 658 | # PyTorch 2.10's FxGraphCachePickler.dumps() doesn't catch ValueError, |
| 659 | # causing torch.compile cache failures when tensors with non-standard |
| 660 | # layouts (e.g. blocked-layout prepacked weights) are serialized. |
| 661 | # PyTorch mainline fixed this in pytorch/pytorch#176557 (merged 2026-03-04). |
| 662 | # This is a thin backport for 2.10 users; remove once 2.10 is dropped. |
| 663 | |
| 664 | |
| 665 | def _apply_fxgraphcache_pickle_patch(pickler_cls, bypass_cls): |
| 666 | """Wrap pickler_cls.dumps to convert ValueError into bypass_cls. |
| 667 | |
| 668 | Idempotent: sets `_vllm_fxgraph_dumps_patched` on the class after the |
| 669 | first apply to prevent re-application. The wrapper function is also |
| 670 | marked with `_vllm_patched` as an additional safeguard. |
| 671 | """ |
| 672 | if getattr(pickler_cls, "_vllm_fxgraph_dumps_patched", False): |
| 673 | return |
| 674 | |
| 675 | original_dumps = pickler_cls.dumps |
| 676 | if hasattr(original_dumps, "_vllm_patched"): |
| 677 | return |
| 678 | |
| 679 | def patched_dumps(self, obj): |
| 680 | try: |
| 681 | return original_dumps(self, obj) |
| 682 | except ValueError as e: |
| 683 | raise bypass_cls("Failed to pickle cache key") from e |
| 684 | |
| 685 | patched_dumps._vllm_patched = True # type: ignore[attr-defined] |
| 686 | pickler_cls.dumps = patched_dumps |
| 687 | pickler_cls._vllm_fxgraph_dumps_patched = True # type: ignore[attr-defined] |
| 688 | |
| 689 | |
| 690 | def _patch_fxgraphcache_pickle_if_needed(): |
| 691 | """Apply FxGraphCachePickler.dumps ValueError backport when on torch 2.10.x.""" |
| 692 | from vllm.utils.torch_utils import is_torch_equal_or_newer |
| 693 | |
| 694 | if not is_torch_equal_or_newer("2.10.0") or is_torch_equal_or_newer("2.11.0"): |
| 695 | return |
| 696 | |
| 697 | from torch._inductor.codecache import BypassFxGraphCache, FxGraphCachePickler |
| 698 | |
| 699 | _apply_fxgraphcache_pickle_patch(FxGraphCachePickler, BypassFxGraphCache) |
| 700 | |
| 701 | |
| 702 | _patch_fxgraphcache_pickle_if_needed() |
| 703 | |
| 704 | # =================================================== |
| 705 | # torch 2.11 Inductor cpp codegen indirect_assert scalar-mask fix |
| 706 | # =================================================== |
| 707 | # CppVecKernel.indirect_assert wraps a scalar mask with |
| 708 | # `VecMask<...>(scalar)`, which is not a valid constructor and triggers a |
| 709 | # C++ compile error during torch.compile of any model that does indirect |
| 710 | # indexing inside a tail-vectorized loop (e.g. Qwen3-VL). |
| 711 | # Failure looks like: |
| 712 | # no matching function for call to 'VecMask<int64_t,2>::VecMask(int&)' |
| 713 | # Upstream fix in PyTorch mainline replaces the call with |
| 714 | # `VecMask<...>::from(scalar)`, see pytorch/pytorch#178148 (lands in 2.12). |
| 715 | # This is a thin backport for torch >= 2.11 and < 2.12; remove once the |
| 716 | # minimum supported torch is 2.12. |
| 717 | |
| 718 | |
| 719 | def _apply_cpp_indirect_assert_patch(): |
| 720 | """Replace CppVecKernel.indirect_assert with a fixed copy that uses |
| 721 | `VecMask<...>::from(scalar)` for scalar masks. |
| 722 | |
| 723 | Idempotent: marks the class with `_vllm_indirect_assert_patched` after |
| 724 | the first apply. |
| 725 | """ |
| 726 | from torch._inductor.codegen.cpp import CppVecKernel |
| 727 | |
| 728 | if getattr(CppVecKernel, "_vllm_indirect_assert_patched", False): |
| 729 | return |
| 730 | |
| 731 | from torch._inductor.codegen.cpp import CppCSEVariable, cexpr_index |
| 732 | |
| 733 | def patched_indirect_assert(self, var, lower, upper, mask=None): |
| 734 | assert isinstance(var, CppCSEVariable) |
| 735 | assert var.dtype is not None |
| 736 | if not var.is_vec: |
| 737 | if isinstance(mask, CppCSEVariable) and mask.is_vec: |
| 738 | mask = f"({mask}).all_masked()" |
| 739 | return super(CppVecKernel, self).indirect_assert(var, lower, upper, mask) |
| 740 | lower_scalar = lower |
| 741 | upper_scalar = upper |
| 742 | if lower: |
| 743 | lower = f"{self._get_vec_type(var.dtype)}({lower})" |
| 744 | if upper: |
| 745 | upper = f"{self._get_vec_type(var.dtype)}({upper})" |
| 746 | if lower and upper: |
| 747 | cond = f"({lower} <= {var}) & ({var} < {upper})" |
| 748 | cond_print = f"{lower_scalar} <= {var} < {upper_scalar}" |
| 749 | elif lower: |
| 750 | cond = f"{lower} <= {var}" |
| 751 | cond_print = f"{lower_scalar} <= {var}" |
| 752 | else: |
| 753 | assert upper |
| 754 | cond = f"{var} < {upper}" |
| 755 | cond_print = f"{var} < {upper_scalar}" |
| 756 | cond = f"{self._get_mask_type(var.dtype)}({cond})" |
| 757 | if mask: |
| 758 | if not mask.is_vec: |
| 759 | # Backport of pytorch/pytorch#178148 -- use ::from for |
| 760 | # scalar masks so g++ picks the correct overload. |
| 761 | mask = f"{self._get_mask_type(var.dtype)}::from({mask})" |
| 762 | cond = f"({cond}) | ~({mask})" |
| 763 | if self.tail_size: |
| 764 | cond = ( |
| 765 | f"{self._get_mask_type(var.dtype)}::set(" |
| 766 | f"{self._get_mask_type(var.dtype)}::from(1)" |
| 767 | f", ({cond}), {cexpr_index(self.tail_size)})" |
| 768 | ) |
| 769 | cond = f"({cond}).all_masked()" |
| 770 | return f'{self.assert_function}({cond}, "index out of bounds: {cond_print}")' |
| 771 | |
| 772 | CppVecKernel.indirect_assert = patched_indirect_assert |
| 773 | CppVecKernel._vllm_indirect_assert_patched = True # type: ignore[attr-defined] |
| 774 | |
| 775 | |
| 776 | def _patch_cpp_indirect_assert_if_needed(): |
| 777 | """Apply cpp codegen indirect_assert backport when on torch 2.11.x. |
| 778 | |
| 779 | Defers application until torch._inductor.codegen.cpp is naturally |
| 780 | imported by Inductor. Importing it eagerly during vllm.__init__ pulls |
| 781 | in torch._inductor.scheduler, whose top-level |
| 782 | `import torch._inductor.async_compile` can fail with |
| 783 | `ModuleNotFoundError: import of torch._inductor.async_compile halted; |
| 784 | None in sys.modules` depending on the import order on the runner |
| 785 | (observed in vLLM CPU CI). |
| 786 | """ |
| 787 | if not is_torch_equal_or_newer("2.11.0") or is_torch_equal_or_newer("2.12.0.dev"): |
| 788 | return |
| 789 | |
| 790 | import sys |
| 791 | |
| 792 | target_name = "torch._inductor.codegen.cpp" |
| 793 | if target_name in sys.modules: |
| 794 | _apply_cpp_indirect_assert_patch() |
| 795 | return |
| 796 | |
| 797 | import importlib.abc |
| 798 | |
| 799 | class _CppCodegenPatchFinder(importlib.abc.MetaPathFinder): |
| 800 | def find_spec(self, fullname, path, target=None): |
| 801 | if fullname != target_name: |
| 802 | return None |
| 803 | sys.meta_path.remove(self) |
| 804 | spec = importlib.util.find_spec(fullname) |
| 805 | if spec is None or spec.loader is None: |
| 806 | return None |
| 807 | original_exec = spec.loader.exec_module |
| 808 | |
| 809 | def _exec_then_patch(module): |
| 810 | original_exec(module) |
| 811 | _apply_cpp_indirect_assert_patch() |
| 812 | |
| 813 | spec.loader.exec_module = _exec_then_patch # type: ignore[method-assign] |
| 814 | return spec |
| 815 | |
| 816 | sys.meta_path.insert(0, _CppCodegenPatchFinder()) |
| 817 | |
| 818 | |
| 819 | _patch_cpp_indirect_assert_if_needed() |
| 820 | |
| 821 | # ============================================================ |
| 822 | # Inductor FALLBACK_ALLOW_LIST fast-path for vllm::*/vllm_aiter::* ops |
| 823 | # ============================================================ |
| 824 | # When Inductor encounters a custom op without a registered lowering or |
| 825 | # decomposition (e.g. vllm::all_reduce, vllm_aiter::fused_add_rms_norm) it |
| 826 | # correctly creates an implicit fallback that calls into the eager Python |
| 827 | # impl. However, unless `base_name` (e.g. "vllm::all_reduce") is in |
| 828 | # torch._inductor.lowering.FALLBACK_ALLOW_LIST, GraphLowering.call_function |
| 829 | # (torch/_inductor/graph.py:~1283) takes the slow path that emits |
| 830 | # log.info("Creating implicit fallback for:\n%s", |
| 831 | # error.operator_str(target, args, kwargs)) |
| 832 | # `operator_str` eagerly recurses through __str__ on every input TensorBox; |
| 833 | # for deep MoE/TP graphs (e.g. Kimi-K2.6 at TP=8) the IR provenance tree |
| 834 | # behind a TP all-reduce input or a residual-fed RMSNorm input is hundreds |
| 835 | # of layers deep, and stringifying it consumes many minutes of CPU per call, |
| 836 | # effectively hanging compilation. |
| 837 | # |
| 838 | # Patching FALLBACK_ALLOW_LIST membership to also match any "vllm::*" or |
| 839 | # "vllm_aiter::*" base_name routes our custom ops through the fast path |
| 840 | # `make_fallback(target, warn=False, override_decomp=True)` instead. This |
| 841 | # preserves all downstream behaviour (allreduce_rms_fusion still pattern- |
| 842 | # matches them, partitioning still works, fallback semantics identical) but |
| 843 | # skips the expensive log formatting on the FIRST encounter of each op. |
| 844 | # |
| 845 | # We wrap the OrderedSet in a thin proxy that: |
| 846 | # - Returns True from __contains__ for any vllm::*/vllm_aiter::* op |
| 847 | # - Otherwise delegates to the underlying set (preserving membership of |
| 848 | # the standard entries like "torchvision::roi_align", "aten::index_add") |
| 849 | # - Forwards add()/__iter__()/__len__()/etc. so other Inductor code paths |
| 850 | # that mutate or iterate the set keep working. |
| 851 | |
| 852 | _VLLM_FALLBACK_NAMESPACE_PREFIXES = ("vllm::", "vllm_aiter::") |
| 853 | |
| 854 | |
| 855 | class _VllmFallbackAllowList: |
| 856 | """Membership proxy that auto-allows vllm::*/vllm_aiter::* base_names.""" |
| 857 | |
| 858 | _vllm_patched = True |
| 859 | |
| 860 | def __init__(self, inner): |
| 861 | self._inner = inner |
| 862 | |
| 863 | def __contains__(self, item): |
| 864 | if isinstance(item, str) and item.startswith(_VLLM_FALLBACK_NAMESPACE_PREFIXES): |
| 865 | return True |
| 866 | return item in self._inner |
| 867 | |
| 868 | def add(self, item): |
| 869 | self._inner.add(item) |
| 870 | |
| 871 | def discard(self, item): |
| 872 | self._inner.discard(item) |
| 873 | |
| 874 | def __iter__(self): |
| 875 | return iter(self._inner) |
| 876 | |
| 877 | def __len__(self): |
| 878 | return len(self._inner) |
| 879 | |
| 880 | def __repr__(self): |
| 881 | return f"_VllmFallbackAllowList({self._inner!r})" |
| 882 | |
| 883 | def __getattr__(self, name): |
| 884 | return getattr(self._inner, name) |
| 885 | |
| 886 | |
| 887 | def _patch_inductor_fallback_allow_list() -> None: |
| 888 | """Wrap torch._inductor.lowering.FALLBACK_ALLOW_LIST so any custom op in |
| 889 | the ``vllm::`` or ``vllm_aiter::`` namespaces is treated as a member. |
| 890 | |
| 891 | Idempotent: a sentinel attribute on the proxy prevents re-wrapping. |
| 892 | """ |
| 893 | try: |
| 894 | from torch._inductor import lowering as _lowering |
| 895 | except ImportError: |
| 896 | return |
| 897 | |
| 898 | base = getattr(_lowering, "FALLBACK_ALLOW_LIST", None) |
| 899 | if base is None or getattr(base, "_vllm_patched", False): |
| 900 | return |
| 901 | |
| 902 | _lowering.FALLBACK_ALLOW_LIST = _VllmFallbackAllowList(base) |
| 903 | |
| 904 | # torch/_inductor/graph.py imports the symbol at module load time: |
| 905 | # from torch._inductor.lowering import FALLBACK_ALLOW_LIST |
| 906 | # so we also need to overwrite the local binding in the graph module if |
| 907 | # it has already been imported. |
| 908 | try: |
| 909 | from torch._inductor import graph as _graph |
| 910 | |
| 911 | if hasattr(_graph, "FALLBACK_ALLOW_LIST"): |
| 912 | _graph.FALLBACK_ALLOW_LIST = _lowering.FALLBACK_ALLOW_LIST |
| 913 | except ImportError: |
| 914 | pass |
| 915 | |
| 916 | |
| 917 | _patch_inductor_fallback_allow_list() |
| 918 | |
| 919 | |
| 920 | def _apply_inductor_pattern_matcher_patch() -> None: |
| 921 | """Allow custom ops and functionalization wrappers with unsupported dtypes.""" |
| 922 | from torch._inductor import pattern_matcher |
| 923 | from torch._inductor.lowering import fallback_node_due_to_unsupported_type |
| 924 | |
| 925 | def fallback_for_builtin(node, allow_cpu_inputs=True): |
| 926 | return ( |
| 927 | isinstance(node.target, torch._ops.OpOverload) |
| 928 | and torch._library.utils.is_builtin(node.target) |
| 929 | and fallback_node_due_to_unsupported_type(node, allow_cpu_inputs) |
| 930 | ) |
| 931 | |
| 932 | pattern_matcher.fallback_node_due_to_unsupported_type = fallback_for_builtin |
| 933 | |
| 934 | |
| 935 | def _patch_inductor_pattern_matcher() -> None: |
| 936 | """Apply the backport when Inductor imports its pattern matcher.""" |
| 937 | import sys |
| 938 | |
| 939 | target_name = "torch._inductor.pattern_matcher" |
| 940 | if target_name in sys.modules: |
| 941 | _apply_inductor_pattern_matcher_patch() |
| 942 | return |
| 943 | |
| 944 | import importlib.abc |
| 945 | |
| 946 | class _PatternMatcherPatchFinder(importlib.abc.MetaPathFinder): |
| 947 | def find_spec(self, fullname, path, target=None): |
| 948 | if fullname != target_name: |
| 949 | return None |
| 950 | sys.meta_path.remove(self) |
| 951 | spec = importlib.util.find_spec(fullname) |
| 952 | if spec is None or spec.loader is None: |
| 953 | return None |
| 954 | original_exec = spec.loader.exec_module |
| 955 | |
| 956 | def _exec_then_patch(module): |
| 957 | original_exec(module) |
| 958 | _apply_inductor_pattern_matcher_patch() |
| 959 | |
| 960 | spec.loader.exec_module = _exec_then_patch # type: ignore[method-assign] |
| 961 | return spec |
| 962 | |
| 963 | sys.meta_path.insert(0, _PatternMatcherPatchFinder()) |
| 964 | |
| 965 | |
| 966 | # Remove once the minimum supported torch includes |
| 967 | # https://github.com/pytorch/pytorch/pull/196013. |
| 968 | if not is_torch_equal_or_newer("2.16.0.dev"): |
| 969 | _patch_inductor_pattern_matcher() |
| 970 | |
| 971 | # ============================================================ |
| 972 | # Triton Autotuner determinism |
| 973 | # ============================================================ |
| 974 | # Replace the Autotuner.run so it always pick the first running configuration. |
| 975 | # Useful to eliminate autotune variability leading to non determinism. |
| 976 | if os.environ.get("VLLM_TRITON_FORCE_FIRST_CONFIG", "0").strip().lower() in ( |
| 977 | "1", |
| 978 | "true", |
| 979 | ): |
| 980 | from vllm.triton_utils.force_first_config import ( |
| 981 | install as _install_force_first_config, |
| 982 | ) |
| 983 | |
| 984 | _install_force_first_config() |