完整部署快照

vllm/env_override.py

@@ -0,0 +1,145 @@
+# SPDX-License-Identifier: Apache-2.0
+import os
+
+import torch
+
+# set some common config/environment variables that should be set
+# for all processes created by vllm and all processes
+# that interact with vllm workers.
+# they are executed whenever `import vllm` is called.
+
+# see https://github.com/NVIDIA/nccl/issues/1234
+os.environ['NCCL_CUMEM_ENABLE'] = '0'
+
+# see https://github.com/vllm-project/vllm/pull/15951
+# it avoids unintentional cuda initialization from torch.cuda.is_available()
+os.environ['PYTORCH_NVML_BASED_CUDA_CHECK'] = '1'
+
+# see https://github.com/vllm-project/vllm/issues/10480
+os.environ['TORCHINDUCTOR_COMPILE_THREADS'] = '1'
+# see https://github.com/vllm-project/vllm/issues/10619
+# torch._inductor.config.compile_threads = 1
+
+# --- BI-V100 Triton compat: inject get_cuda_stream into triton.runtime.jit ---
+# torch._inductor.triton_heuristics expects triton.runtime.jit.get_cuda_stream
+# but Triton 2.3.1 (standard, BI-V100 image) does not export it.
+# Upstream PyTorch >=2.5 switched to torch._C._cuda_getCurrentRawStream.
+# We inject a compat shim so the import succeeds in worker processes too.
+# This must run before any torch._inductor import.
+def _patch_triton_get_cuda_stream():
+    try:
+        import triton.runtime.jit as _jit
+    except ImportError:
+        return
+    if hasattr(_jit, 'get_cuda_stream'):
+        return
+    def _get_cuda_stream(device_index: int = 0) -> int:
+        if hasattr(torch._C, '_cuda_getCurrentRawStream'):
+            return torch._C._cuda_getCurrentRawStream(device_index)
+        return 0
+    _jit.get_cuda_stream = _get_cuda_stream
+
+_patch_triton_get_cuda_stream()
+del _patch_triton_get_cuda_stream
+
+# --- BI-V100 ixformer compat: bridge SDK 0.6.0 infer API to CoreX 3.2.3 ---
+# SDK 0.6.0 inference/functions/*.py calls ops.infer.xxx but CoreX 3.2.3
+# _C.so (v0.3.0) only has _C._functions.xxx_forward. We bridge them.
+# Must run before any ixformer.inference.functions import.
+def _patch_ixformer_infer():
+    try:
+        import ixformer._C as _C
+    except ImportError:
+        return
+    if hasattr(_C, 'infer'):
+        return
+    try:
+        # Use the full bridge module if deployed
+        import importlib.util, os
+        shim_path = os.path.join(os.path.dirname(__file__),
+                                 '..', '..', 'patch_ixformer_infer.py')
+        if not os.path.exists(shim_path):
+            # Try qwen3_6_scripts location
+            for candidate in [
+                '/usr/local/corex/lib64/python3/dist-packages/vllm/patch_ixformer_infer.py',
+                os.path.join(os.path.dirname(__file__), 'patch_ixformer_infer.py'),
+            ]:
+                if os.path.exists(candidate):
+                    shim_path = candidate
+                    break
+        if os.path.exists(shim_path):
+            spec = importlib.util.spec_from_file_location(
+                'patch_ixformer_infer', shim_path)
+            mod = importlib.util.module_from_spec(spec)
+            spec.loader.exec_module(mod)
+            return
+    except Exception:
+        pass
+
+    # Minimal inline fallback: just map linear
+    if not hasattr(_C, '_functions'):
+        return
+    import types
+    infer = types.ModuleType('ixformer._C.infer')
+    _fn = _C._functions
+
+    def _linear(input, weight, act_type=-1, bias=None, output=None,
+                persistent=False):
+        import torch.nn.functional as F
+        result = F.linear(input, weight, bias)
+        if act_type == 3: result = F.gelu(result)
+        elif act_type == 4: result = F.relu(result)
+        elif act_type == 12: result = F.silu(result)
+        if output is not None:
+            output.copy_(result)
+            return output
+        return result
+
+    infer.linear = _linear
+    infer.linear_ex = lambda inp, w, b=None, o=None: _linear(inp, w, -1, b, o)
+    infer.mixed_type_linear = lambda inp, w, b=None, o=None: _linear(
+        inp.to(w.dtype), w, -1, b, o)
+    _C.infer = infer
+
+_patch_ixformer_infer()
+del _patch_ixformer_infer
+
+# --- BI-V100 ixformer.distributed: inject missing gather/send/recv ---
+# SDK 0.6.0 distributed has gather() but CoreX 3.2.3 system install doesn't.
+# Inject it using all_gather_into_tensor (which does exist).
+def _patch_ixformer_distributed():
+    try:
+        import ixformer.distributed as ixfd
+    except ImportError:
+        return
+    if hasattr(ixfd, 'gather'):
+        return
+
+    def gather(tensor, gather_list=None, dst=0, group=None,
+               async_op=False, use_comm_stream=False):
+        """gather via all_gather_into_tensor + slice on non-dst ranks."""
+        import torch
+        import torch.distributed as dist
+        world_size = ixfd.get_group_world_size(group) if group else 1
+        if world_size <= 1:
+            if gather_list is not None:
+                gather_list[0].copy_(tensor)
+            return
+        # Use all_gather_into_tensor which exists in this version
+        flat_output = torch.empty(
+            [world_size] + list(tensor.shape),
+            dtype=tensor.dtype, device=tensor.device)
+        ixfd.all_gather_into_tensor(
+            flat_output.view(-1), tensor.contiguous(),
+            group=group, async_op=async_op)
+        # Get current rank via torch.distributed (reliable across versions)
+        rank = dist.get_rank()
+        if rank == dst and gather_list is not None:
+            for i in range(world_size):
+                gather_list[i].copy_(flat_output[i])
+
+    ixfd.gather = gather
+    print('[ixformer_compat] injected gather into ixformer.distributed')
+
+_patch_ixformer_distributed()
+del _patch_ixformer_distributed