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

.github/configs/cuda.yml

@@ -85,6 +85,14 @@ integration_tests:
     command: >-
       python -m pytest tests/integration/test_profiler_parity.py
       -v -m main_ops --tb=short
+  # Both compile files run in one pytest invocation on purpose. The compile
+  # 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 how the autotune and
+  # worker-import failures stayed hidden until a multi-file run reproduced them.
+  - name: Run torch.compile tests
+    command: >-
+      python -m pytest tests/integration/test_compile.py
+      tests/integration/test_compile_autotune.py -v -s --tb=short
   - name: Run inference tests
     command: >-
       python -m pytest tests/integration/test_qwen3_infer.py

docs/torch_compile_integration.md

@@ -143,11 +143,27 @@ stream-less autograd nodes and AOT autograd's backward trace trips
 
 ## Performance
 
-**Not yet measured.** Correctness is verified (`tests/integration/test_compile.py`);
-benchmarking the fusion gain, and comparing it against stock `inductor` on cuda,
-is still open work. Structurally the two should land close together -- same
-inductor fusion passes, same Triton codegen, and since the graph stays on flagos
-there is no per-call copy -- but that is an expectation, not a measurement.
+Measured on one A100 (torch 2.10 CPU wheel + external cu128 `libtorch_cuda.so`),
+fp32, comparing compiled against eager on flagos:
+
+| Workload | Eager | Compiled | Speedup |
+|---|---|---|---|
+| Qwen3-0.6B forward, 1x128 tokens | 35.6 ms | 15.9 ms | **2.24x** |
+| Elementwise chain, 4096x4096 | 0.97 ms | 0.11 ms | **9.18x** |
+| Elementwise chain, 1024x1024 | 0.086 ms | 0.051 ms | 1.69x |
+| Transformer block (`bench_compile.py --model=transformer`) | 1.41 ms | 1.00 ms | 1.41x |
+| MLP 2048x4096 (matmul-bound) | 8.43 ms | 7.96 ms | 1.06x |
+| MLP 64x512 (too small to amortize launch) | 0.18 ms | 0.19 ms | 0.92x |
+
+The pattern is what inductor's fusion predicts: the win comes from collapsing
+elementwise chains into one kernel, so it scales with how much of the graph is
+elementwise and how large the tensors are. Matmul-dominated graphs still call
+cuBLAS and barely move, and at small sizes the launch overhead of the compiled
+wrapper can exceed the saving.
+
+Not yet measured: a like-for-like comparison against stock `inductor` on cuda
+(`--compare-cuda` exists but has not been run), and any training/backward
+throughput number.
 
 ### Benchmarking
 
@@ -245,10 +261,18 @@ tests live alongside it:
 4. **CUDA graphs off**: `torch.cuda.CUDAGraph` is a dummy class in the CPU torch
    wheel, so `triton.cudagraphs` is forced off even under `mode="max-autotune"`
 5. **FlagTree maturity**: Backend support varies by hardware (NVIDIA most mature)
+6. **Convolutions do not compile**: inductor prefers `channels_last` for conv on
+   GPU, and while the flagos convolution kernel honours that layout, its
+   fake/meta kernel still predicts contiguous strides -- so inductor rejects the
+   graph on a stride mismatch (`aten.convolution.default`). Eager never hits
+   this, since it is the layout pass that produces a `channels_last` input.
+   Reproduce with `python tests/perf/bench_compile.py --model=conv`.
 
 ## Roadmap
 
 - [x] Phase 1: Inductor integration (flagos as a first-class GPU device)
+- [x] Measure fusion gains against eager on flagos (see Performance)
+- [ ] Fix the conv `channels_last` meta/real stride mismatch (Limitations #6)
 - [ ] Phase 2: FlagTree integration — shim exists, not yet exercised end-to-end
 - [ ] Benchmark fusion gains vs. stock inductor+triton on cuda
 - [ ] Phase 3: FlagGems-aware fusion (recognize pre-optimized patterns)

tests/integration/test_compile_autotune.py

@@ -0,0 +1,169 @@
+# Copyright 2026 FlagOS Contributors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Regression tests for the two gaps that only bite *multi-kernel* graphs.
+
+``test_compile.py`` compiles single-``Linear`` models, which need neither
+Triton autotuning nor parallel compilation, so both of the failures guarded
+here slipped through it:
+
+1. **Autotuning needs a constructible event.** Inductor's
+   ``InductorBenchmarker.get_event_pairs`` builds ``torch.cuda.Event(
+   enable_timing=True)``. On the CPU-torch wheel that class derives from a
+   dummy type and raises on construction, so any graph with more than one
+   Triton config to choose between died with "Tried to instantiate dummy base
+   class Event".
+
+2. **Compile workers need torch_fl.** Inductor farms kernel compilation out to
+   subprocesses that import only ``torch`` and ``triton``. Without ``torch_fl``
+   the CUDA driver never registers, and triton reports "Could not find an
+   active GPU backend".
+
+Both are properties of the *shape* of the graph rather than of any op, so the
+models below are chosen to force multiple kernels: stacked linears (autotuning)
+plus normalizations and reductions (enough kernels to go parallel).
+"""
+
+import pytest
+import torch
+import torch_fl
+
+
+@pytest.fixture
+def device():
+    if torch_fl.flagos.device_count() == 0:
+        pytest.skip("No flagos devices available")
+    return "flagos:0"
+
+
+@pytest.fixture(autouse=True)
+def _fresh_dynamo():
+    """Each test compiles from scratch; a cached graph would prove nothing."""
+    torch._dynamo.reset()
+    yield
+    torch._dynamo.reset()
+
+
+def _assert_matches_eager(model, *args, rtol=1e-3, atol=1e-3):
+    """Compile ``model``, and check it lands on flagos with eager's numbers."""
+    expected = model(*args)
+    actual = torch.compile(model, backend="flagos")(*args)
+    assert actual.device.type in ("privateuseone", "flagos"), (
+        f"output landed on {actual.device}, expected flagos"
+    )
+    torch.testing.assert_close(actual, expected, rtol=rtol, atol=atol)
+    return actual
+
+
+def test_event_is_constructible(device):
+    """The autotuner's exact call. Guards gap 1 at its narrowest point."""
+    start = torch.cuda.Event(enable_timing=True)
+    end = torch.cuda.Event(enable_timing=True)
+
+    x = torch.randn(2048, 2048, device=device)
+    start.record()
+    for _ in range(10):
+        x = x * 1.001
+    end.record()
+    end.synchronize()
+
+    # Real device timing, not a host stand-in: work this size is never free,
+    # and never takes a minute.
+    elapsed = start.elapsed_time(end)
+    assert 0.0 < elapsed < 60_000.0, f"implausible elapsed_time: {elapsed}"
+
+
+def test_flagos_event_is_constructible(device):
+    """``flagos.Event`` is what ``torch.cuda.Event`` is routed to."""
+    event = torch_fl.flagos.Event(enable_timing=True)
+    event.record()
+    event.synchronize()
+    assert event.query() is True
+
+
+def test_compile_stacked_linears(device):
+    """Two matmuls give the autotuner configs to pick between (gap 1)."""
+    model = torch.nn.Sequential(
+        torch.nn.Linear(512, 512),
+        torch.nn.GELU(),
+        torch.nn.Linear(512, 512),
+    ).to(device)
+    _assert_matches_eager(model, torch.randn(64, 512, device=device))
+
+
+def test_compile_normalizations(device):
+    """LayerNorms emit enough kernels to reach parallel compilation (gap 2)."""
+    model = torch.nn.Sequential(
+        torch.nn.LayerNorm(512),
+        torch.nn.Linear(512, 512),
+        torch.nn.GELU(),
+        torch.nn.Linear(512, 512),
+        torch.nn.LayerNorm(512),
+    ).to(device)
+    _assert_matches_eager(model, torch.randn(64, 512, device=device))
+
+
+def test_compile_reductions(device):
+    """Softmax/sum/mean are reduction kernels -- a different codegen path."""
+
+    def f(x):
+        return (x * 2.0).softmax(-1).sum(-1) + x.mean(-1)
+
+    _assert_matches_eager(f, torch.randn(256, 512, device=device))
+
+
+def test_compile_backward_multikernel(device):
+    """Backward doubles the kernel count, so it hits both gaps hardest."""
+    model = torch.nn.Sequential(
+        torch.nn.Linear(512, 512),
+        torch.nn.GELU(),
+        torch.nn.Linear(512, 512),
+    ).to(device)
+    x = torch.randn(64, 512, device=device, requires_grad=True)
+
+    torch.compile(model, backend="flagos")(x).sum().backward()
+
+    assert x.grad is not None, "no gradient reached the input"
+    assert x.grad.device.type in ("privateuseone", "flagos")
+    assert torch.isfinite(x.grad).all(), "gradient has non-finite entries"
+
+
+def test_compile_dynamic_shapes(device):
+    """One compile serving several batch sizes still has to autotune."""
+    model = torch.nn.Sequential(
+        torch.nn.Linear(512, 512),
+        torch.nn.ReLU(),
+        torch.nn.Linear(512, 512),
+    ).to(device)
+    compiled = torch.compile(model, backend="flagos", dynamic=True)
+
+    for batch in (16, 32, 64):
+        x = torch.randn(batch, 512, device=device)
+        actual = compiled(x)
+        assert actual.shape == (batch, 512)
+        torch.testing.assert_close(actual, model(x), rtol=1e-3, atol=1e-3)
+
+
+def test_compile_max_autotune(device):
+    """``max-autotune`` benchmarks candidates, so it cannot skip the events."""
+    model = torch.nn.Sequential(
+        torch.nn.Linear(256, 256),
+        torch.nn.GELU(),
+        torch.nn.Linear(256, 256),
+    ).to(device)
+    x = torch.randn(64, 256, device=device)
+
+    expected = model(x)
+    actual = torch.compile(model, backend="flagos", mode="max-autotune")(x)
+    torch.testing.assert_close(actual, expected, rtol=1e-3, atol=1e-3)

tests/perf/bench_compile.py

@@ -25,10 +25,17 @@
 """
 
 import argparse
+import os
 import time
+
+# torch_fl must be imported before torch: it preloads the external
+# libtorch_cuda.so that this build's `torch.cuda` depends on. Import torch
+# first and the very first flagos tensor dies with "Cannot initialize CUDA
+# without ATen_cuda library". The test suite gets away without ordering its
+# own imports because conftest imports torch_fl during collection.
+import torch_fl
 import torch
 import torch.nn as nn
-import torch_fl
 
 
 class MLPModel(nn.Module):
@@ -48,7 +55,7 @@ def forward(self, x):
         x = x + 1.0
 
         x = self.fc2(x)
-        x = torch.gelu(x)
+        x = nn.functional.gelu(x)
         x = x / 2.0
 
         x = self.fc3(x)
@@ -115,11 +122,17 @@ def benchmark_model(model, inputs, warmup=10, rounds=100):
         with torch.no_grad():
             _ = model(*inputs) if isinstance(inputs, tuple) else model(inputs)
 
+    # flagos tensors report either name depending on how the device was spelled,
+    # so check both -- matching only "privateuseone" would silently time a
+    # "flagos:0" model through torch.cuda.synchronize().
+    def sync():
+        if device.type in ("privateuseone", "flagos"):
+            torch_fl.flagos.synchronize()
+        else:
+            torch.cuda.synchronize()
+
     # Sync before timing
-    if device.type == "privateuseone":
-        torch_fl.flagos.synchronize()
-    else:
-        torch.cuda.synchronize()
+    sync()
 
     # Timed runs
     start = time.perf_counter()
@@ -128,10 +141,7 @@ def benchmark_model(model, inputs, warmup=10, rounds=100):
             _ = model(*inputs) if isinstance(inputs, tuple) else model(inputs)
 
     # Sync after timing
-    if device.type == "privateuseone":
-        torch_fl.flagos.synchronize()
-    else:
-        torch.cuda.synchronize()
+    sync()
 
     elapsed = time.perf_counter() - start
     return (elapsed / rounds) * 1000  # Convert to ms
@@ -182,9 +192,7 @@ def main():
     print(f"Batch size: {args.batch_size}")
     print(f"Hidden size: {args.hidden_size}")
     print(f"Rounds: {args.rounds}")
-    print(
-        f"FlagTree enabled: {bool(int(torch.os.environ.get('FLAGOS_USE_FLAGTREE', '0')))}"
-    )
+    print(f"FlagTree enabled: {os.environ.get('FLAGOS_USE_FLAGTREE', '0') == '1'}")
     print()
 
     # === Flagos Device ===

torch_fl/compile/README.md

@@ -76,10 +76,49 @@ backend compensates:
   that raises on construction; `mode="max-autotune"` would otherwise enable it.
 - `CudaInterface.get_raw_stream` is re-attached -- the binding exists, but the
   import-time `torch.cuda._is_compiled()` probe left it at `None`.
+- `worker_start_method = "fork"` -- see below.
 
 See `torch_fl/accelerator/cuda/_cuda_compat.py` for the memory-stats and
 Event/Stream shims that inductor's autotuner needs.
 
+### Autotuning and compile workers
+
+Two failures only appear once a graph is big enough to need more than one Triton
+kernel, which is why the original single-`Linear` tests missed both.
+`tests/integration/test_compile_autotune.py` guards them.
+
+**Autotuning needs a constructible event.** `InductorBenchmarker.get_event_pairs`
+times candidate configs with `torch.cuda.Event(enable_timing=True)`. In the CPU
+wheel that class derives from a `torch._utils._dummy_type` placeholder and raises
+on construction. `torch_fl.flagos.Event` therefore switches base class: on a
+vendor torch build it still subclasses `torch.cuda.Event`, but when that is a
+dummy it subclasses the device-agnostic `torch.Event`, which dispatches to our
+own `c10::flagos::DeviceGuardImpl` (`csrc/runtime/guard.h`) for
+record/block/query/elapsedTime. 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 patch has to land on `torch.cuda.Event`; 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 start out able to see the device --
+keeping compilation parallel, unlike `compile_threads = 1` (Qwen3-0.6B: 31.9s
+forked vs 40.8s serial). Both overrides are scoped to the flagos build by probing
+for a missing `torch._C` CUDA binding, so a vendor torch install keeps inductor's
+defaults.
+
+Because the worker pool is created lazily and shared, a single test can be served
+before the pool exists. Run the two compile test files together (as
+`.github/configs/cuda.yml` does) or these failures can hide.
+
 ## Environment Variables
 
 - `FLAGOS_USE_FLAGTREE=1` - Reserved for FlagTree integration (Phase 2)
@@ -89,9 +128,14 @@ Event/Stream shims that inductor's autotuner needs.
 
 1. Single device - multi-GPU compilation not yet exercised
 2. FlagTree integration is a stub (Phase 2)
+3. Convolutions do not compile: the flagos conv kernel honours `channels_last`
+   but its fake/meta kernel predicts contiguous strides, so inductor's conv
+   layout pass produces a graph it then rejects on a stride mismatch. See
+   `docs/torch_compile_integration.md` (Limitations).
 
 ## Future Work
 
+- [ ] Fix the conv `channels_last` meta/real stride mismatch
 - [ ] FlagTree integration to replace OpenAI Triton
 - [ ] Benchmark fusion gains against stock inductor+triton on cuda
 - [ ] Multi-GPU compilation support