PPU async Inductor workers reinitialize CUDA after fork #1031

Summary

FlagTree's PPU backend fails in PyTorch Inductor's asynchronous compile workers after the parent process has initialized CUDA. The worker calls PPUDriver.get_active_torch_device(), which calls torch.cuda.current_device() and attempts to initialize CUDA in a forked subprocess.

Serial compilation with TORCHINDUCTOR_COMPILE_THREADS=1 avoids the worker process and makes the full Torch-FL compile suite pass.

Environment

• Hardware: T-Head PPU, 16 visible devices
• FlagTree commit: 52678a8b781d2a3e36fc1a77fe74a9aa23c5ca4b
• FlagTree package: 0.6.0+ppu.git52678a8b
• Triton API version: 3.6.0
• Build: FLAGTREE_BACKEND=ppu MAX_JOBS=16 pip install . --no-build-isolation
• PyTorch: 2.10.0, CUDA compatibility version 13.0
• Consumer: Torch-FL 2e00b393cf80088706b460a187aef185d3a283f4
• Python: 3.12.13

FlagTree detection is correct:

triton 3.6.0
triton._flagtree_spec present
FLAGTREE_BACKEND ppu

Reproduction

Run Torch-FL's compile suite with the default Inductor worker count (32 in this environment) and a fresh cache:

PYTHONPATH=/tmp/flagtree-ppu-target \
FLAGOS_DISABLE_CUDA_ASSETS=1 \
FLAGOS_USE_FLAGTREE=1 \
TORCHINDUCTOR_CACHE_DIR=/tmp/torchinductor-flagtree-ppu-default \
pytest tests/integration/test_compile.py -v --tb=short

Observed result:

13 passed, 2 failed

FAILED tests/integration/test_compile.py::test_compile_backward
FAILED tests/integration/test_compile.py::test_compile_dtypes[dtype1]

Both failures end with:

File "triton/compiler/hint_manager.py", line 127, in hint_get_flagtree_backend
  device = driver.active.get_active_torch_device()
File "triton/backends/ppu/driver.py", line 682, in get_active_torch_device
  return torch.device("cuda", self.get_current_device())
File "torch/cuda/__init__.py", line 1094, in current_device
  _lazy_init()
RuntimeError: Cannot re-initialize CUDA in forked subprocess. To use CUDA with
multiprocessing, you must use the 'spawn' start method

The failure is reproducible with a clean Inductor cache. Which tests fail depends on when a new worker-side compilation is needed, but backward and FP16 consistently exercise it in the full suite.

Workaround and control results

Force in-process/serial compilation:

PYTHONPATH=/tmp/flagtree-ppu-target \
FLAGOS_DISABLE_CUDA_ASSETS=1 \
FLAGOS_USE_FLAGTREE=1 \
TORCHINDUCTOR_COMPILE_THREADS=1 \
pytest tests/integration/test_compile.py -v --tb=short

Result:

15 passed, 14 warnings in 5.22s

The dedicated FlagTree end-to-end test also passes:

tests/integration/test_compile.py::test_flagtree_compiles_correct_results PASSED

A compiled PPU smoke test matches eager and returns on flagos:0:

triton 3.6.0
flagtree_active True
flagtree_backend ppu
device flagos:0
match True

This indicates that PPU lowering, code generation, launch, FP16, and backward are functional; the failure is specific to device discovery in the asynchronous compile worker.

Root cause

triton/compiler/hint_manager.py::hint_get_flagtree_backend() asks the active driver for a torch device. The PPU implementation is:

def get_active_torch_device(self):
    import torch
    return torch.device("cuda", self.get_current_device())

self.get_current_device() is torch.cuda.current_device. In a worker created after CUDA initialization in the parent, this triggers PyTorch's fork-safety guard.

The PPU driver should avoid CUDA runtime initialization when the hint manager only needs a device/backend hint, or the compile worker must use a CUDA-safe process start/initialization path.

Expected behavior

FlagTree PPU compilation should work with Inductor's default asynchronous compilation configuration after the parent has initialized a PPU/CUDA context. Calling get_active_torch_device() in the worker should not attempt to reinitialize CUDA after fork.

Possible directions

• Make the PPU hint path derive the backend/device without calling torch.cuda.current_device() in the worker.
• Pass the selected device index from the parent into compile workers.
• Ensure the relevant worker uses a CUDA-safe start method.
• Add a regression test that initializes CUDA in the parent and compiles a new PPU kernel asynchronously.