Skip to content

fix(compile): make autotuning and parallel kernel compilation work on flagos - #59

Open
lvyufeng wants to merge 1 commit into
flagos-ai:mainfrom
lvyufeng:fix/compile-autotune-and-workers
Open

fix(compile): make autotuning and parallel kernel compilation work on flagos#59
lvyufeng wants to merge 1 commit into
flagos-ai:mainfrom
lvyufeng:fix/compile-autotune-and-workers

Conversation

@lvyufeng

@lvyufeng lvyufeng commented Aug 6, 2026

Copy link
Copy Markdown
Collaborator

The torch.compile integration merged in #41 only ever compiled single-Linear
models in its tests, which need neither autotuning nor more than one Triton
kernel. Two independent failures hid behind that. Both reproduce on any graph
with a couple of stacked Linears or a LayerNorm.

Autotuning needs a constructible event

InductorBenchmarker.get_event_pairs times candidate configs with
torch.cuda.Event(enable_timing=True). In the CPU-only wheel this build pairs
with an external libtorch_cuda.so, that binding was never compiled, so
torch.cuda substitutes a placeholder from torch._utils._dummy_type whose
__new__ raises "Tried to instantiate dummy base class Event".
flagos.Event subclassed it and inherited the failure.

flagos.Event now picks its base class by lineage: on a vendor torch build it
still subclasses torch.cuda.Event, and when that is a dummy it subclasses the
device-agnostic torch.Event, which dispatches record/block/query/elapsedTime
to c10::flagos::DeviceGuardImpl (csrc/runtime/guard.h). Timing stays a real
device measurement, and since every vendor under csrc/runtime/accelerator/
implements that ABI, the fallback is portable rather than NVIDIA-specific.

The fix has to land here. Patching triton.testing.do_bench does not help,
because inductor reaches the benchmarker through
triton_heuristics.benchmark_all_configs -> bench -> benchmarker.benchmark_gpu,
not through do_bench.

Compile workers need torch_fl

Inductor's default worker_start_method, "subprocess", starts workers as a
bare sys.executable -m torch._inductor.compile_worker that imports only torch
and triton. flagos lives behind PrivateUse1, so such a worker has no
accelerator: triton's CudaDriver.is_active() asks torch.cuda.is_available(),
gets False, and the worker dies with "Could not find an active GPU backend".

"fork" inherits this process, torch_fl included, so workers come up already
seeing the device — and compilation stays parallel, unlike compile_threads = 1
(Qwen3-0.6B: 31.9s forked vs 40.8s serial). Both overrides are scoped to this
build by probing for a missing torch._C CUDA binding, so a vendor torch
install keeps inductor's defaults.

Testing

tests/integration/test_compile_autotune.py guards both: stacked Linears,
normalizations, reductions, multi-kernel backward, dynamic shapes and
max-autotune, plus a direct check that the autotuner's own Event call works.
On a cleared TORCHINDUCTOR_CACHE_DIR, 7 of its 8 tests fail before these
changes.

CI ran no compile tests at all, so .github/configs/cuda.yml now runs both
compile files — in one pytest invocation on purpose. The worker pool is created
lazily and shared, so a file run on its own can be served entirely before the
pool spins up, which is exactly how the worker failure stayed hidden until a
multi-file run reproduced it.

Verified on one A100 (torch 2.10 CPU wheel + external cu128 libtorch_cuda.so):

  • Both compile files, cold cache, three separate runs: 20 passed, 1 skipped
  • tests/integration/ops -m "main_ops and not flaggems_python and not flaggems_cpp": 117 passed, 15 skipped, 1 xfailed
  • Unit tests: 15 passed
  • ruff check / ruff format --check clean on all touched files
  • Fork's multi-threaded DeprecationWarning stress-tested over 6 sequential
    compile/pool cycles, no hangs; spawn was also tried and fails (1 of 4)

Performance

Measured fp32, compiled vs eager on flagos:

Workload Eager Compiled Speedup
Qwen3-0.6B forward, 1x128 35.6 ms 15.9 ms 2.24x
Elementwise chain 4096x4096 0.97 ms 0.11 ms 9.18x
Transformer block 1.41 ms 1.00 ms 1.41x
Elementwise chain 1024x1024 0.086 ms 0.051 ms 1.69x
MLP 2048x4096 (matmul-bound) 8.43 ms 7.96 ms 1.06x
MLP 64x512 0.18 ms 0.19 ms 0.92x

Qwen3 numerics match eager at rtol/atol 2e-2. The pattern is what inductor's
fusion predicts: the win comes from collapsing elementwise chains, matmul-bound
graphs still call cuBLAS, and at small sizes launch overhead exceeds the saving.

Also in here

tests/perf/bench_compile.py had never been run and could not be: it called
torch.gelu (nonexistent), read torch.os.environ, recognised only the
"privateuseone" spelling of the device, and imported torch before torch_fl —
which the docs now state as a hard requirement, since torch_fl preloads the
libtorch_cuda.so that torch.cuda depends on. The test suite gets away
without ordering its imports because conftest imports torch_fl during
collection.

Known limitation documented, not fixed

Convolutions do not compile. Inductor prefers channels_last for conv on GPU,
and while the flagos conv kernel honours that layout, its fake/meta kernel
still predicts contiguous strides, so inductor rejects the graph on a stride
mismatch in aten.convolution.default. Confirmed by direct comparison on a
16x3x224x224 channels_last input — flagos real (3211264, 1, 14336, 64) vs
fake (3211264, 50176, 224, 1), while cuda agrees with itself. Eager never
hits it, since it is the layout pass that produces a channels_last input.
Reproduce with bench_compile.py --model=conv. Filed as Limitations #6.

… flagos

The torch.compile integration merged in flagos-ai#41 only ever compiled single-Linear
models in its tests, which need neither autotuning nor more than one Triton
kernel. Two independent failures hid behind that. Both reproduce on any
graph with a couple of stacked Linears or a LayerNorm.

Autotuning needs a constructible event. InductorBenchmarker.get_event_pairs
times candidate configs with torch.cuda.Event(enable_timing=True). In the
CPU-only wheel this build pairs with an external libtorch_cuda.so, that
binding was never compiled, so torch.cuda substitutes a placeholder from
torch._utils._dummy_type whose __new__ raises "Tried to instantiate dummy
base class Event". flagos.Event subclassed it and inherited the failure.

flagos.Event now picks its base class by lineage: on a vendor torch build it
still subclasses torch.cuda.Event, and when that is a dummy it subclasses
the device-agnostic torch.Event, which dispatches record/block/query/
elapsedTime to c10::flagos::DeviceGuardImpl (csrc/runtime/guard.h). Timing
stays a real device measurement, and since every vendor under
csrc/runtime/accelerator/ implements that ABI, the fallback is portable
rather than NVIDIA-specific. Note the fix has to land here: patching
triton.testing.do_bench does not help, because inductor reaches the
benchmarker through triton_heuristics.benchmark_all_configs -> bench ->
benchmarker.benchmark_gpu, not through do_bench.

Compile workers need torch_fl. Inductor's default worker_start_method,
"subprocess", starts workers as a bare `sys.executable -m
torch._inductor.compile_worker` that imports only torch and triton. flagos
lives behind PrivateUse1, so such a worker has no accelerator: triton's
CudaDriver.is_active() asks torch.cuda.is_available(), gets False, and the
worker dies with "Could not find an active GPU backend". "fork" inherits
this process, torch_fl included, so workers come up already seeing the
device -- and compilation stays parallel, unlike compile_threads = 1
(Qwen3-0.6B: 31.9s forked vs 40.8s serial). Both overrides are scoped to
this build by probing for a missing torch._C CUDA binding, so a vendor
torch install keeps inductor's defaults.

tests/integration/test_compile_autotune.py guards both: stacked Linears,
normalizations, reductions, multi-kernel backward, dynamic shapes and
max-autotune, plus a direct check that the autotuner's own Event call works.
On a cleared TORCHINDUCTOR_CACHE_DIR, 7 of its 8 tests fail before these
changes.

CI ran no compile tests at all, so .github/configs/cuda.yml now runs both
compile files -- in one pytest invocation on purpose. The worker pool is
created lazily and shared, so a file run on its own can be served entirely
before the pool spins up, which is exactly how the worker failure stayed
hidden until a multi-file run reproduced it.

Measured on one A100, fp32, compiled vs eager on flagos: Qwen3-0.6B forward
2.24x (35.6ms -> 15.9ms, numerics matching eager at rtol/atol 2e-2),
elementwise chain 4096x4096 9.18x, transformer block 1.41x, matmul-bound
MLP 1.06x, and 0.92x at 64x512 where launch overhead exceeds the saving.

tests/perf/bench_compile.py had never been run and could not be: it called
torch.gelu (nonexistent), read torch.os.environ, recognised only the
"privateuseone" spelling of the device, and imported torch before torch_fl
-- which the docs now state as a hard requirement, since torch_fl preloads
the libtorch_cuda.so that torch.cuda depends on. The test suite gets away
without it because conftest imports torch_fl during collection.

Also documents a third bug found while benchmarking and left unfixed:
convolutions do not compile. Inductor prefers channels_last for conv on
GPU, and while the flagos conv kernel honours that layout, its fake/meta
kernel still predicts contiguous strides, so inductor rejects the graph on
a stride mismatch. Eager never hits it, since it is the layout pass that
produces a channels_last input. Reproduce with bench_compile.py --model=conv.
@lvyufeng
lvyufeng force-pushed the fix/compile-autotune-and-workers branch from 9dd33f4 to e375d3e Compare August 6, 2026 02:58
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants