From e375d3e1b69aca35a1ef069eecc06c4132b86745 Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 6 Aug 2026 02:29:39 +0000 Subject: [PATCH] fix(compile): make autotuning and parallel kernel compilation work on flagos 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. 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. --- .github/configs/cuda.yml | 8 + docs/torch_compile_integration.md | 34 ++++- tests/integration/test_compile_autotune.py | 169 +++++++++++++++++++++ tests/perf/bench_compile.py | 34 +++-- torch_fl/compile/README.md | 44 ++++++ torch_fl/compile/inductor_backend.py | 20 +++ torch_fl/flagos/__init__.py | 116 ++++++++++---- 7 files changed, 375 insertions(+), 50 deletions(-) create mode 100644 tests/integration/test_compile_autotune.py diff --git a/.github/configs/cuda.yml b/.github/configs/cuda.yml index ff6f089d..50256c17 100644 --- a/.github/configs/cuda.yml +++ b/.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 diff --git a/docs/torch_compile_integration.md b/docs/torch_compile_integration.md index 3b47ee5a..04f9b242 100644 --- a/docs/torch_compile_integration.md +++ b/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) diff --git a/tests/integration/test_compile_autotune.py b/tests/integration/test_compile_autotune.py new file mode 100644 index 00000000..ba63ad52 --- /dev/null +++ b/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) diff --git a/tests/perf/bench_compile.py b/tests/perf/bench_compile.py index 9a2dbabc..c777cdbc 100644 --- a/tests/perf/bench_compile.py +++ b/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 === diff --git a/torch_fl/compile/README.md b/torch_fl/compile/README.md index a8bd00ca..da7cb5c1 100644 --- a/torch_fl/compile/README.md +++ b/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 diff --git a/torch_fl/compile/inductor_backend.py b/torch_fl/compile/inductor_backend.py index 778aea54..77631df4 100644 --- a/torch_fl/compile/inductor_backend.py +++ b/torch_fl/compile/inductor_backend.py @@ -101,6 +101,26 @@ def _resolve_config_patches( if not hasattr(torch._C, "_StaticCudaLauncher"): patches["use_static_cuda_launcher"] = False + # Kernel compilation is farmed out to worker processes. Inductor's default, + # "subprocess", launches them as a bare `sys.executable -m + # torch._inductor.compile_worker` (subproc_pool.py), which imports only + # torch and triton -- never torch_fl. flagos lives behind PrivateUse1, so a + # worker without it has no accelerator registered: triton's + # `CudaDriver.is_active()` consults `torch.cuda.is_available()`, gets False, + # and the worker dies with "Could not find an active GPU backend". Any graph + # with enough kernels to go parallel fails. + # + # "fork" inherits this process, torch_fl included, so the workers come up + # already able to see the device. That keeps compilation parallel -- forcing + # `compile_threads = 1` would also work, but measurably slower (Qwen3-0.6B: + # 31.9s forked vs 40.8s serial). + # + # Only applies to the flagos build described above. A vendor torch install + # registers CUDA in every fresh interpreter, so its workers are fine as + # spawned and we leave inductor's default alone. + if not hasattr(torch._C, "_cuda_getDevice"): + patches.setdefault("worker_start_method", "fork") + return patches diff --git a/torch_fl/flagos/__init__.py b/torch_fl/flagos/__init__.py index a47906cf..d13adc7b 100644 --- a/torch_fl/flagos/__init__.py +++ b/torch_fl/flagos/__init__.py @@ -284,40 +284,92 @@ def _real_current_stream(device=None): ) -class Event(torch.cuda.Event): - """Flagos event that wraps a CUDA event (same physical GPU under boxing). - - Since flagos shares the CUDA stream/GPU, a real ``torch.cuda.Event`` gives - true device-side semantics (maca event under the hood): ``record`` and - ``wait`` enforce cross-stream ordering, ``elapsed_time`` measures on-device - time, and ``query`` reflects actual completion -- unlike the previous - host-timestamp stand-in whose ``wait`` was a no-op. - - ``record``/``wait`` are overridden to default to the *real* current stream - (see :func:`_real_current_stream`) rather than ``torch.cuda.current_stream``, - which boxing replaces with a non-Stream shim. +def _cuda_event_is_usable(): + """Is ``torch.cuda.Event`` backed by a real binding, or a dummy stand-in? + + On a vendor torch build the base class is the C++ ``_CudaEventBase``. On the + CPU-only wheel this project pairs with an external ``libtorch_cuda.so``, the + binding was never compiled, so ``torch.cuda`` substitutes a placeholder + built by ``torch._utils._dummy_type`` -- a bare ``object`` subclass whose + ``__new__`` raises "Tried to instantiate dummy base class Event". Detect it + by that lineage rather than by trying to construct one, since a failed + construction is indistinguishable from a genuine device error. """ + base = torch.cuda.Event.__mro__[-2] # the class just above `object` + return base.__module__ != "torch._utils" + + +if _cuda_event_is_usable(): + + class Event(torch.cuda.Event): + """Flagos event backed by a real CUDA event (same physical GPU). + + Since flagos shares the CUDA stream/GPU, a real ``torch.cuda.Event`` + gives true device-side semantics (maca event under the hood): ``record`` + and ``wait`` enforce cross-stream ordering, ``elapsed_time`` measures + on-device time, and ``query`` reflects actual completion -- unlike the + previous host-timestamp stand-in whose ``wait`` was a no-op. + + ``record``/``wait`` default to the *real* current stream (see + :func:`_real_current_stream`) rather than ``torch.cuda.current_stream``, + which boxing replaces with a non-Stream shim. + """ + + def __new__( + cls, enable_timing=False, blocking=False, interprocess=False, external=False + ): + return super().__new__( + cls, + enable_timing=enable_timing, + blocking=blocking, + interprocess=interprocess, + external=external, + ) - def __new__( - cls, enable_timing=False, blocking=False, interprocess=False, external=False - ): - return super().__new__( - cls, - enable_timing=enable_timing, - blocking=blocking, - interprocess=interprocess, - external=external, - ) - - def record(self, stream=None): - if stream is None: - stream = _real_current_stream() - return super().record(stream) - - def wait(self, stream=None): - if stream is None: - stream = _real_current_stream() - return super().wait(stream) + def record(self, stream=None): + if stream is None: + stream = _real_current_stream() + return super().record(stream) + + def wait(self, stream=None): + if stream is None: + stream = _real_current_stream() + return super().wait(stream) + +else: + + class Event(torch.Event): + """Flagos event on the device-agnostic ``torch.Event``. + + Used when ``torch.cuda.Event`` is a dummy (the CPU-wheel build). The + generic event dispatches straight to our own + ``c10::flagos::DeviceGuardImpl`` -- ``record``/``block``/``queryEvent``/ + ``elapsedTime`` in ``csrc/runtime/guard.h`` -- so timing is measured by + the vendor runtime, exactly as the cuda-backed branch above. Every + vendor in ``csrc/runtime/accelerator/`` implements that ABI, so this is + portable rather than NVIDIA-specific. + + This is what makes inductor's autotuner work: it constructs + ``torch.cuda.Event(enable_timing=True)`` to time candidate Triton + configs, and ``_cuda_compat`` points that name here. + + ``record``/``wait`` take no stream: flagos submits everything to the + default stream, and the generic event already defaults to the current + one, so there is no shim to route around. + """ + + def __new__( + cls, enable_timing=False, blocking=False, interprocess=False, external=False + ): + # `external` has no equivalent on torch.Event and is unused here + # (nothing in this project imports an event from another library). + return super().__new__( + cls, + device=f"flagos:{current_device()}", + enable_timing=enable_timing, + blocking=blocking, + interprocess=interprocess, + ) def current_stream(device=None):