145 lines
4.9 KiB
Python
145 lines
4.9 KiB
Python
"""
|
|
BI-V100 Triton/torch._inductor compatibility patch.
|
|
|
|
Problem:
|
|
torch._inductor.triton_heuristics (corex torch build) line 43:
|
|
from triton.runtime.jit import get_cuda_stream, KernelInterface
|
|
But Triton 2.3.1 (standard, shipped in BI-V100 image) does not export
|
|
get_cuda_stream from triton.runtime.jit.
|
|
|
|
Root cause:
|
|
PyTorch's inductor was written against an older Triton that re-exported
|
|
get_cuda_stream. Newer Triton moved it out; upstream PyTorch >=2.5
|
|
imports from torch._C._cuda_getCurrentRawStream instead, and wraps
|
|
Triton symbols through torch._inductor.runtime.triton_compat.
|
|
|
|
Fix:
|
|
1. Inject get_cuda_stream into triton.runtime.jit at the Python level
|
|
so the old torch._inductor import succeeds.
|
|
2. Optionally patch the .py file on disk (deploy-time) so the fix
|
|
survives across process forks and reimports.
|
|
|
|
This module should be imported BEFORE torch._inductor is loaded.
|
|
Safe to import multiple times (idempotent).
|
|
"""
|
|
|
|
import importlib
|
|
import sys
|
|
|
|
|
|
def _get_cuda_stream_from_torch(device_index: int = 0) -> int:
|
|
"""Get the current CUDA stream pointer, matching the old Triton API."""
|
|
import torch
|
|
if hasattr(torch._C, '_cuda_getCurrentRawStream'):
|
|
return torch._C._cuda_getCurrentRawStream(device_index)
|
|
# Fallback: return default stream (0)
|
|
return 0
|
|
|
|
|
|
def patch_triton_runtime_jit():
|
|
"""
|
|
Ensure triton.runtime.jit.get_cuda_stream exists.
|
|
|
|
If triton.runtime.jit is already imported and lacks get_cuda_stream,
|
|
inject it. If not yet imported, inject after import. This makes
|
|
torch._inductor.triton_heuristics importable.
|
|
"""
|
|
try:
|
|
import triton.runtime.jit as jit_module
|
|
except ImportError:
|
|
# Triton not installed at all; nothing to patch
|
|
return False
|
|
|
|
if hasattr(jit_module, 'get_cuda_stream'):
|
|
# Already present (compatible Triton version), nothing to do
|
|
return False
|
|
|
|
# Inject the compat shim
|
|
jit_module.get_cuda_stream = _get_cuda_stream_from_torch
|
|
return True
|
|
|
|
|
|
def patch_triton_heuristics_on_disk(torch_inductor_path: str = None):
|
|
"""
|
|
Patch the torch._inductor triton_heuristics.py file on disk to use
|
|
torch._C._cuda_getCurrentRawStream instead of triton.runtime.jit.get_cuda_stream.
|
|
|
|
This is the deploy-time fix applied by patch_ops.sh.
|
|
|
|
Args:
|
|
torch_inductor_path: Path to torch/_inductor directory.
|
|
If None, auto-detected from torch installation.
|
|
|
|
Returns:
|
|
True if file was patched, False if already patched or not found.
|
|
"""
|
|
import os
|
|
|
|
if torch_inductor_path is None:
|
|
import torch
|
|
torch_root = os.path.dirname(torch.__file__)
|
|
torch_inductor_path = os.path.join(torch_root, '_inductor')
|
|
|
|
# Check both old location and new location
|
|
candidates = [
|
|
os.path.join(torch_inductor_path, 'triton_heuristics.py'),
|
|
os.path.join(torch_inductor_path, 'runtime', 'triton_heuristics.py'),
|
|
]
|
|
|
|
patched = False
|
|
for filepath in candidates:
|
|
if not os.path.exists(filepath):
|
|
continue
|
|
|
|
with open(filepath, 'r') as f:
|
|
content = f.read()
|
|
|
|
# Idempotency: skip if already patched (try/except block present)
|
|
if '_cuda_getCurrentRawStream as get_cuda_stream' in content:
|
|
print(f'[skip] already patched: {filepath}')
|
|
continue
|
|
|
|
old_import = 'from triton.runtime.jit import get_cuda_stream'
|
|
if old_import not in content:
|
|
continue
|
|
|
|
# Replace the problematic import
|
|
# Handle the case where KernelInterface is on the same line
|
|
if 'from triton.runtime.jit import get_cuda_stream, KernelInterface' in content:
|
|
old_line = 'from triton.runtime.jit import get_cuda_stream, KernelInterface'
|
|
new_block = """try:
|
|
from triton.runtime.jit import get_cuda_stream, KernelInterface
|
|
except ImportError:
|
|
from torch._C import _cuda_getCurrentRawStream as get_cuda_stream
|
|
try:
|
|
from triton.runtime.jit import KernelInterface
|
|
except ImportError:
|
|
from triton.compiler import KernelInterface"""
|
|
content = content.replace(old_line, new_block)
|
|
else:
|
|
new_import = """try:
|
|
from triton.runtime.jit import get_cuda_stream
|
|
except ImportError:
|
|
from torch._C import _cuda_getCurrentRawStream as get_cuda_stream"""
|
|
content = content.replace(old_import, new_import)
|
|
|
|
with open(filepath, 'w') as f:
|
|
f.write(content)
|
|
|
|
# Clear __pycache__ for the patched file
|
|
cache_dir = os.path.join(os.path.dirname(filepath), '__pycache__')
|
|
if os.path.isdir(cache_dir):
|
|
import shutil
|
|
shutil.rmtree(cache_dir, ignore_errors=True)
|
|
|
|
patched = True
|
|
print(f'[ok] patched triton_heuristics: {filepath}')
|
|
|
|
return patched
|
|
|
|
|
|
# Auto-apply runtime patch on import
|
|
_applied = patch_triton_runtime_jit()
|
|
if _applied:
|
|
print('[triton_compat] injected get_cuda_stream into triton.runtime.jit')
|