This document describes the torch.compile integration for the flagos device, enabling automatic kernel fusion and optimization via TorchInductor.
The flagos device now supports PyTorch 2.0+ torch.compile for automatic performance optimization:
import torch_fl # Import first on MetaX.
import torch
model = MyModel().to("flagos:0")
compiled_model = torch.compile(model, backend="flagos")
# Automatic fusion of elementwise ops, reduced dispatch overhead
output = compiled_model(input)Key benefits:
- Automatic kernel fusion (no manual optimization needed), cutting per-op dispatch overhead
- Graph stays on flagos: no cuda round trip, no copy at the graph boundary
- Compatible with existing flagos dispatch (FlagGems Python/C++, CUDA boxing)
- Optional FlagTree compilation for multi-backend kernel generation (drop-in: it
replaces
tritonat install time, so no code change is needed)
import torch_fl # Import first on MetaX.
import torch
# Standard model definition
model = torch.nn.Sequential(
torch.nn.Linear(512, 512),
torch.nn.ReLU(),
torch.nn.Linear(512, 512),
).to("flagos:0")
# Compile with flagos backend
model = torch.compile(model, backend="flagos")
# Use as normal
x = torch.randn(64, 512, device="flagos:0")
y = model(x) # Automatically fused kernels# Default: balanced optimization
model = torch.compile(model, backend="flagos")
# Maximum performance (longer compile, better runtime)
model = torch.compile(model, backend="flagos", mode="max-autotune")
# Explicit inductor config overrides
model = torch.compile(model, backend="flagos", options={"max_autotune": True})mode and options are expanded into inductor config patches scoped to that
compile. Note that CUDA graphs are always forced off (see Limitations), so
mode="reduce-overhead" -- whose main lever is cudagraphs -- has little effect
here.
FlagTree is a Triton fork whose kernel compiler targets many vendor backends. It integrates by substitution at install time, which is the whole thing to understand about it:
- Its wheel is named
flagtree, but the module it installs istriton. - Installing it uninstalls the official
tritonand takes its place. - So inductor's own
import tritonalready resolves to FlagTree once installed, and nothing intorch_flpatchessys.modules.
Wheels from 0.6.2 on also install a real flagtree package, which is the
FlagPrism debugger/profiler host -- it borrows the DSL out of
triton.language.core rather than providing a compiler of its own. Importing it
is still not how you reach FlagTree; reach for triton and check which build it
is with is_flagtree_active().
There is no flagtree package on PyPI; build it from source. On a machine whose
triton is in use by FlagGems, build into a separate virtualenv, since the
install removes the existing triton:
# System deps (nlohmann-json must match cmake/json-version.txt)
apt install zlib1g-dev libxml2-dev nlohmann-json3-dev
git clone https://github.com/flagos-ai/FlagTree.git
cd FlagTree
pip install -r python/requirements.txt
# Pick the vendor backend. Do NOT set this for nvidia or amd.
# export FLAGTREE_BACKEND=metax
# Repo root for Triton 3.4+ (branch main); use `cd python` first on 3.1-3.3.
MAX_JOBS=64 pip install . --no-build-isolation -vLLVM arrives as a prebuilt tarball (~4.5 GB unpacked), so this does not build
LLVM from source. Verify, from any directory other than FlagTree/python:
pip show flagtree # wheel name
python -c 'import triton; print(triton.__path__)' # module nameBranch selects the Triton base version, which must match what your torch
expects: main is Triton 3.6 (correct for torch 2.10), triton_v3.5.x is 3.5
and carries the Ascend backend.
Since FlagTree's install removes the existing triton, the cheapest way to test
against it is to build into a separate virtualenv and shadow triton by path,
keeping torch and torch_fl from your main environment. Triton does not link
against torch, so the two mix freely as long as the Python versions match:
PYTHONPATH=/path/to/flagtree-venv/lib/python3.12/site-packages \
FLAGOS_USE_FLAGTREE=1 pytest tests/integration/test_compile.pyThis avoids rebuilding torch_fl inside the FlagTree venv, and leaves the
FlagGems-facing triton in the main environment untouched.
Once built, compilation goes through FlagTree with no further configuration.
Setting FLAGOS_USE_FLAGTREE=1 asserts that the active triton is FlagTree
and errors out if it is stock Triton — it cannot switch FlagTree on, because that
was decided when the environment was built:
FLAGOS_USE_FLAGTREE=1 python your_script.pyDetection uses the FlagTree-only module triton._flagtree_spec. Note that
triton._flagtree_backend.FLAGTREE_BACKEND is not a usable signal: it is the
empty string on nvidia/amd builds, since upstream directs you not to set
FLAGTREE_BACKEND for those.
PPU FlagTree currently queries torch.cuda.current_device() while selecting
compiler hints. Inductor's asynchronous compile workers can make that query
after a PPU CUDA context was initialized in the parent, which PyTorch rejects as
CUDA reinitialization after fork (FlagTree #1031).
The flagos backend therefore defaults PPU FlagTree to one compile thread until
the upstream driver no longer initializes CUDA in a worker. An explicit
TORCHINDUCTOR_COMPILE_THREADS value or compile_threads compile option remains
authoritative.
MUSA uses the vendor FlagTree runtime (flagtree-0.5.0+mthreads3.1, Triton 3.1)
and its mthreads Python backend. The compiler emits a musa target. A stock
Triton wheel only provides nvidia/amd backends and cannot compile MUSA
kernels.
FlagTree talks to torch_fl directly; the separate torch_musa plugin is not
required and its __init__ is never imported. The vendor MThreads driver reads
four runtime facts from torch_musa — device availability (which is what makes
Triton select the backend at all), current device, compute capability, and the
raw musaStream_t. Running that plugin's __init__ in this process is not an
option, because it claims the process-global PrivateUse1 hooks that torch_fl
must own. So
torch_fl/compile/flagtree_shim.py
answers those four questions from torch_fl's own MUSA runtime and rebinds them
onto the vendor driver class before the first driver instance exists. The
consequence that matters: compiled kernels are submitted to the same stream the
native mudnn kernels use, so a compiled graph is ordered against eager MUSA work
without an explicit synchronize between them.
The binding must happen before Triton instantiates its driver, since the driver's
__init__ copies these attributes onto the instance. flagos_compile_backend
therefore binds at module load and again per compile, both idempotently.
tests/integration/test_compile.py::test_musa_flagtree_binds_to_torch_fl_runtime
asserts the binding took effect by checking that every rebound driver callable
resolves inside torch_fl. That, rather than torch_musa being absent from
sys.modules, is the property worth pinning: the driver would also "work"
against some other runtime, and torch_fl itself publishes a small
compatibility surface under that module name when FlagGems is enabled (see
_install_musa_flaggems_compat).
MThreads FlagTree currently queries the MUSA runtime while compiling. The flagos
backend therefore defaults this path to one Inductor compile thread, preserving
an explicit TORCHINDUCTOR_COMPILE_THREADS or compile_threads setting. This
serialization is a safety workaround for the vendor driver, not a performance
claim; remove it only after the relevant FlagTree driver is verified fork-safe.
On the measured MTT S5000 setup, validation requires the vendor FlagTree runtime,
FLAGOS_ACCELERATOR=musa, and FLAGOS_USE_FLAGTREE=1. The generic Triton 3.7.1 runtime
is not MThreads execution evidence.
Example environment:
PYTHONPATH=/path/to/flagtree-mthreads-runtime:$PWD \
LD_LIBRARY_PATH=/path/to/flagtree-mthreads-runtime/triton/_C:/usr/local/musa/lib:$LD_LIBRARY_PATH \
TORCH_DEVICE_BACKEND_AUTOLOAD=0 FLAGOS_ACCELERATOR=musa FLAGOS_USE_FLAGTREE=1 \
pytest tests/integration/test_compile.py -vThe focused FlagTree test must compare compiled output with eager output and
assert that outputs and gradients remain on flagos; CPU-only tests can cover
registration and vendor target selection but do not establish MUSA compiler
support.
Ascend's FlagGems route now runs on FlagTree (0.6.2a1+ascend3.5, Triton
3.5), which installs itself as the triton package and registers an
AscendBackend. FLAGOS_USE_FLAGTREE=1 asserts that the active triton really
is FlagTree — it cannot switch anything on, because the wheel replaces triton
at install time — so on a FlagTree environment
test_flagtree_compiles_correct_results runs rather than skipping.
The compile path described below was measured against triton-ascend 3.2.0,
the toolchain Ascend used before this move, and has not been re-measured on
FlagTree; no CI step exercises it either way. Treat the numbers and the three
workarounds as unvalidated on the current toolchain. The rest of this section
explains why a torch_npu-free policy is needed at all, and that part is
toolchain-independent.
FlagTree's Ascend backend routes its host runtime through torch_npu, which
claims PrivateUse1 on import, after which torch_fl cannot register flagos —
the same conflict documented under vendor setup. This is a genuine
incompatibility with the plugin model, not a packaging detail.
The coupling is deeper than the import torch_npu lines suggest
(driver.py:231, utils.py:48, backend_register.py:87). FlagTree dispatches
these operations through a backend policy, and the torch_npu policy also
decides generated C++ and link flags: get_cc_cmd emits -ltorch_npu,
header_file includes <torch_npu/csrc/core/npu/NPUWorkspaceAllocator.h>, and
allocate_sync_block_lock/async_launch emit at_npu::native:: calls
(backend_register.py:227, :293, :316, :340). So stubbing sys.modules
cannot fix it — the launcher .so needs real headers and libtorch_npu.so.
Upstream's own allocate_memory allocates on at::kPrivateUse1
(backend_register.py:305), i.e. it wants exactly the slot flagos needs.
The seam that does work is the registry itself: register(category, method)
accepts any category string, and mindspore already exists as a second,
non-torch_npu policy. torch_fl/compile/flagtree_ascend_policy.py registers a
third, flagos, backed by torch.flagos and the ACL stream registry, emitting
plain ATen against PrivateUse1. It has to set utils.backend_policy directly,
because get_backend_func only honours TRITON_BACKEND when it names
torch_npu or mindspore (utils.py:42-44), and it forces
TRITON_ENABLE_TASKQUEUE=false since the task queue is torch_npu-only.
Upstream request to make this unnecessary:
flagos-ai/FlagTree#1046
Verified against a real FlagTree build (triton_v3.5.x @ d2063b06,
flagtree-0.6.0+ascend, built with the prebuilt LLVM 7d5de303 and
TRITON_CODEGEN_BACKENDS=nvidia;amd;ascend on aarch64/Python 3.10, and again
against the flagtree-0.6.2a1+ascend3.5 wheel the Ascend FlagGems route now
runs on): after
install_policy(), the registry reports ['flagos', 'mindspore', 'torch_npu'],
flagos covers all 15 required strategies with no parity gap against
torch_npu, and each signature matches upstream's. The generated output is
clean: header_file emits <ATen/ATen.h> with no torch_npu/at_npu
reference, and get_cc_cmd does not link -ltorch_npu. torch_fl.flagos
calls install_policy() during device init, ahead of the first FlagGems
kernel, so eagerly launched kernels take the flagos policy too.
The build also exposed a second, import-order-dependent coupling, separate from
the policy: backends/ascend/__init__.py:24 imports do_bench_npu, and
testing.py:28 imports torch_npu at module scope. Only profiling needs it,
but it runs during backend discovery, before any policy can be selected. With
torch_fl holding PrivateUse1, a real torch_npu refuses to load
(Two accelerators cannot be used at the same time). In practice this is
already absorbed: torch_fl/__init__.py:828 installs a torch_npu stub for
FlagGems' sake, and that stub satisfies FlagTree's import too — verified on the
real build, where import torch_fl first yields backend_policy = flagos with
PrivateUse1 named flagos. The failure only appears if triton is imported
before torch_fl, which the stub-ordering guard already exists to prevent.
Making the upstream import lazy would remove the ordering constraint
altogether; install_policy() reports this case with an actionable message
rather than letting the bare "npu and npu" error through.
The naming side, at least, is already compatible: FlagTree's Ascend backend
reports name.conf = ascend and supports_target accepts only
target.backend == "npu" — exactly the pair _ASCEND_PROFILE already encodes,
so triton_device_type/triton_backend_key need no change. The three
workarounds below would still need re-testing, since that backend is separate
code with its own defect surface.
No extra configuration is needed beyond
the environment FlagGems already requires
(vendor setup) — torch.compile(backend="flagos")
picks the Ascend profile from FLAGOS_ACCELERATOR=ascend.
Measured on a real 910 (Ascend910_9382, CANN 9.0.0, triton-ascend 3.2.0,
torch 2.10.0+cpu, Python 3.10): forward, backward, fused elementwise, and
matmul+normalization graphs all compile and match eager. This measurement is
stale: it was taken on triton-ascend, and the toolchain Ascend ships is now
FlagTree. Nothing has re-run the compile suite on FlagTree, so whether these
numbers, and whether the three workarounds below, still hold is unmeasured.
Support is
experimental — three vendor-toolchain defects had to be worked around, and
each workaround is a place where a toolchain upgrade should let us delete code:
| Defect | Symptom | Workaround |
|---|---|---|
| Masked 2-D load of an 8-bit dtype is miscompiled | Silently reads only the first two elements of each row, repeated. Surfaced as wrong relu gradients (4090/4096 elements off, max abs diff 3.05), no error raised |
triton_byte_loads.py: pass enable_linearize=True on every compile, and rewrite bool kernel args from *i1 to *i8. Both are needed — linearize does not fix the *i1 pointer type, and the identical bytes read correctly through *i8 |
ub overflow is raised as a generic compilation error |
An oversized autotune config fails the whole compile instead of being dropped | triton_resource_limits.py: parse requires N bits while M bits available and re-raise as Triton's OutOfResources, which inductor already knows to skip |
| A kernel built in a compile worker segfaults when the parent launches it | Crash in NPULauncher.__call__ (triton/backends/ascend/driver.py), only with a cold inductor cache and two or more compiles in one process |
inductor_backend.py defaults Ascend to compile_threads=1. fork and spawn both crash, so this is not the CUDA-after-fork problem the PPU path above hits |
triton_libdevice.py additionally fills the Ascend backend's libdevice module
map, which the vendor backend leaves empty.
The first row's enable_linearize requirement is exercised directly by
test_ascend_masked_byte_load_requires_linearize, which asserts the bug is
still present without the option. That test failing is the signal that the
toolchain was fixed and the workaround can go.
flagos is registered with TorchInductor as a first-class GPU device. The
traced graph is handed to compile_fx unchanged -- still on flagos -- and
inductor generates Triton kernels that operate on flagos tensors directly.
Components:
-
Backend registration (
torch_fl/compile/inductor_backend.py)- Registers
"flagos"withtorch._dynamo.register_backend - Expands
mode/optionsinto inductorconfig_patches - Delegates to
compile_fxwith no graph rewriting
- Registers
-
Platform profile (
torch_fl/compile/platform_profile.py)- flagos has no Triton backend of its own, so the compile path must name the
hardware's backend. The profile carries the three facts that differ per
vendor: the name reported as
DeviceProperties.type(and therefore asGPUTarget.backend), the key that backend is registered under intriton.backends.backends, and whether the runtime is CUDA-like at all - The two names are not the same string everywhere. On Ascend the driver
reports
GPUTarget(backend="npu", ...)andAscendBackend.supports_targetaccepts only"npu", but the package key is"ascend" is_cuda_like=False(Ascend only) routes hardware queries, raw streams and generated device snippets away fromtorch.cuda, which does not exist there
- flagos has no Triton backend of its own, so the compile path must name the
hardware's backend. The profile carries the three facts that differ per
vendor: the name reported as
-
Device interface (
torch_fl/compile/device_interface.py)DeviceInterfacesubclass: device state always fromtorch.flagos; hardware properties fromtorch.cudaon CUDA-like builds (the same physical GPU), fromtorch.flagosplus the ACL runtime on Ascend- Adds
"flagos"to inductor'sGPU_TYPESsois_gpu()is True and the Triton codegen path is taken instead of C++/CPU - Reports the hardware's backend name at the Triton boundary
(
DeviceProperties.create), because each vendor backend hard-checkstarget.backend. Target and package key per vendor:cuda/nvidiaon CUDA-compatible builds,maca/metaxon MetaX,npu/ascendon Ascend - On Ascend,
get_compute_capabilityreturns the SoC name (e.g.Ascend910_9382) rather than a number, because that is what reaches Triton asGPUTarget.arch
-
Codegen registration (
torch_fl/compile/inductor_codegen.py)- CUDA-like builds: device op overrides inheriting the CUDA ones, and the stock CUDA/Triton scheduling + wrapper pipeline
- Ascend: overrides that emit the ACL raw stream and
torch.flagosdevice calls, plainTritonScheduling, and no C++ wrapper —CppWrapperGpuemits CUDA-runtime C++, so the slot is left unregistered rather than pointing at something that cannot build. The C++ members raiseNotImplementedErrorinstead of emitting a translation unit CANN rejects
-
Dispatch integration
- Ops inductor does not fuse fall back to eager flagos dispatch (FlagGems Python/C++ or CUDA boxing) with no changes needed
-
GCU 64-bit guard (
torch_fl/compile/triton_64bit_guard.py)- GCU300 cannot compile a kernel that needs a 64-bit value, and one of the
two ways to reach that state segfaults below inductor's own error handling.
Wraps
CachingAutotuner._precompile_configto raise a namedInductorErrorinstead, beforetriton.compileis entered - GCU-only (
platform_profile().vendor, checked inside the patch), so other targets are untouched. Installed from bothflagos_compile_backendandtorch_fl's eager registration, since the defaultbackend="inductor"path never reaches the former
- GCU300 cannot compile a kernel that needs a 64-bit value, and one of the
two ways to reach that state segfaults below inductor's own error handling.
Wraps
Flow:
torch.compile(model, backend="flagos")
→ dynamo captures FX graph (on flagos)
→ compile_fx / AOT autograd, graph never leaves flagos
→ inductor generates fused Triton kernels for flagos tensors
→ unfused ops fall back to flagos eager dispatch
Why the graph is not rewritten to cuda. An earlier version converted the
graph and example inputs to cuda first. Beyond the copy per call, this breaks
backward: at::getAccelerator() is PrivateUse1/flagos, and
torch::autograd::Node::stream() only yields a stream when a node's input
device type equals the accelerator, so a cuda-rewritten graph produces
stream-less autograd nodes and AOT autograd's backward trace trips
opt_ready_stream && opt_parent_stream (engine.cpp:1085).
Components:
-
Detection only (
torch_fl/compile/flagtree_shim.py)- No import patching: FlagTree is the
tritonmodule once installed, so inductor picks it up with no interception is_flagtree_active()tests for the FlagTree-onlytriton._flagtree_specFLAGOS_USE_FLAGTREE=1callsrequire_flagtree(), which errors if the active Triton is stock
- No import patching: FlagTree is the
-
Backend selection
- Chosen at FlagTree build time via
FLAGTREE_BACKEND(unset for nvidia/amd), not at runtime - Same Triton kernel code, different backend compiler
- Chosen at FlagTree build time via
Benefits:
- Multi-backend: same compiled model runs on NVIDIA/Ascend/Cambricon/MetaX
- Future-proof for non-NVIDIA hardware
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.
# Run performance benchmark
python tests/perf/bench_compile.py --model=mlp --batch-size=64
# Compare with CUDA baseline
python tests/perf/bench_compile.py --model=transformer --compare-cuda
# Test FlagTree integration
FLAGOS_USE_FLAGTREE=1 python tests/perf/bench_compile.py| Variable | Default | Description |
|---|---|---|
FLAGOS_USE_FLAGTREE |
0 |
Require the active triton to be FlagTree (assert, not switch) |
FLAGOS_COMPILE_FALLBACK_EAGER |
0 |
Fall back to eager on compile errors |
Existing dispatch variables (FLAGOS_BACKEND_CONFIG) still apply to compiled kernels.
Symptom: torch.compile raises errors during graph capture or codegen.
Solutions:
- Enable fallback to eager:
FLAGOS_COMPILE_FALLBACK_EAGER=1 - Check for unsupported ops (dynamic shapes, custom ops)
- Verify meta implementations for custom ops
Symptom: a compile on GCU fails with
InductorError: RuntimeError: The GCU300 target has no 64-bit support, and kernel
'triton_red_fused_...' needs 64-bit values:
line 21: tmp1 = tl.full([1, 1], 7, tl.int64)
The GCU300 target has no 64-bit support at all, in either of the two forms a generated kernel can reach it:
- a 64-bit operand (an int64 or float64 tensor), visible in the kernel
signature as
*i64/*fp64; - a 64-bit promotion of size arithmetic, which keeps an int32 signature.
Inductor's
SIMDScheduling.select_index_dtypefalls back totorch.int64whenevercan_use_32bit_indexingdeclines the buffer set, and every size-derived value in the kernel then carries that dtype. No input tensor is 64-bit in this case, and no inductor config knob selects the index dtype.
The second form is why torch_fl/compile/triton_64bit_guard.py refuses the
kernel before triton.compile is entered rather than translating an error
afterwards. The vendor compiler cannot legalize the arith.extsi that feeds an
i64 value and aborts the pass pipeline: the process dies with SIGSEGV, exit code
139, no traceback. That abort is below the try/except inductor wraps around
triton.compile in CachingAutotuner._precompile_config, so nothing on
inductor's exception path runs and no Python-level handler can catch it. The
guard is a predicate over the generated source and the signature, and it raises
in place of the crash.
What to do:
- Keep the values 32-bit where the model allows it.
- Skip compiling the kernel that generates it. Flex attention's
create_block_maskhits the promotion on any input dtype, so pass_compile=False(or wrap the factory intorch._dynamo.disable) — the mask builds eagerly and the rest of the model still compiles.torch_fl.compile.triton_64bit_guard.gcu_64bit_unsupported()answers True on such a build, for code that wants to choose the route itself. FLAGOS_COMPILE_FALLBACK_EAGER=1turns the refusal into an eager run (backend="flagos"only — the guard also covers the defaultbackend="inductor"path, which has no fallback wrapper).
This is a limitation of the target, not of the guard: with the guard removed the same kernels come back as a pipeline error or a segfault. A build whose vendor compiler gains 64-bit support should drop the guard; it is gated on the GCU vendor string, so it cannot affect any other platform.
Symptom: Compiled model runs at same speed as eager.
Possible causes:
- Model is compute-bound (large matmuls) — fusion won't help much
- Compilation didn't fuse ops (check inductor logs)
- Dispatch overhead is small relative to kernel time
Debug: Run with TORCH_LOGS="+inductor" to see fusion decisions.
Symptom: FLAGOS_USE_FLAGTREE=1 raises "the active 'triton' module is stock
Triton, not FlagTree".
The env's triton is the official one. FlagTree is selected when the environment
is built, so this cannot be fixed at runtime — build FlagTree per the section
above, into this interpreter. Check what you have with:
python -c 'import triton; print(triton.__path__)'
python -c 'import importlib.util as u; print(u.find_spec("triton._flagtree_spec") is not None)'Do not reach for import flagtree when debugging this; on recent wheels it
imports, but what it gives you is the FlagPrism debugger/profiler namespace, not
the compiler. The compiler is triton, whichever build it is.
# Run integration tests
pytest tests/integration/test_compile.py -v
# Test specific scenarios
pytest tests/integration/test_compile.py::test_basic_compile
pytest tests/integration/test_compile.py::test_compile_backward
# Regression guards for the two codegen fixes this integration required
pytest tests/integration/test_compile.py -k fake_tensor
pytest tests/integration/ops/test_clamp_dispatch.py -v
# Test FlagTree compilation (requires a FlagTree-built env)
FLAGOS_USE_FLAGTREE=1 pytest tests/integration/test_compile.py::test_flagtree_compiles_correct_results
# Platform-profile selection and the vendor workarounds, on any platform
# including plain CPU -- these test the selection logic, not the toolchain
pytest tests/unit/test_compile_platform_profile.py -v
# The 64-bit predicate and the guard's install/refuse behaviour, on any
# platform; the GCU arm of it needs the vendor Triton stack and a real card
pytest tests/unit/test_triton_64bit_guard.py -v
pytest tests/integration/test_compile.py -k "64bit or block_mask" -v
# Ascend, on a real 910. Clear the inductor cache first: a warm cache hides the
# compile-worker crash the serial-compile default exists for
rm -rf /tmp/torchinductor_root
TORCH_DEVICE_BACKEND_AUTOLOAD=0 pytest tests/integration/test_compile.py -vTwo generated-kernel bugs only surface under compilation, so their regression tests live alongside it:
detachre-dispatch. The generated kernel calledat::detach(self), which is registered on PrivateUse1 too and so dispatched back into itself. Eager hid the recursion becauseDeviceBoxingGuardrewrites self's device metadata first; under FakeTensor it cannot, since the Python dispatch key sits above the backend key. Dynamo traces everynn.Linearthrough detach, so this was a stack-overflow segfault at trace time. Fixed by emittingat::native::detach(NATIVE_DIRECT_VIEW_OPSinscripts/codegen/codegen_ops.py).optional<Tensor>boxing in in-place kernels.gen_inplacehanded only plainat::Tensorargs toDeviceBoxingGuard, soclamp_.Tensorpassed unboxed flagosmin/maxinto a CUDAselfand crashed. Fixed by materializing each optional into a holder, matchinggen_functional_pure.
- torch >= 2.0 required: Older PyTorch versions don't have
torch.compile - Inductor-compatible ops only: Custom C++ ops may not fuse
- Dynamic shapes: Some models with dynamic shapes may not compile
- CUDA graphs off:
torch.cuda.CUDAGraphis a dummy class in the CPU torch wheel, sotriton.cudagraphsis forced off even undermode="max-autotune" - FlagTree maturity: Backend support varies by hardware; Hygon HCU is
validated on
gfx936and MetaX on C550/MACA 3.8.0, while other vendor backends remain untested here - Ascend is experimental: compiles serially by default, has no C++ wrapper
codegen (
CppWrapperGpuemits CUDA-runtime C++), and carries three toolchain workarounds — see Ascend above. Those workarounds were measured againsttriton-ascend 3.2.0and have not been re-tested on FlagTree, the toolchain the FlagGems route now runs on. Only the graphs intests/integration/test_compile.pyare validated; whole-model compilation is not yet exercised there - GCU has no 64-bit kernels: the GCU300 target cannot compile a kernel that needs a 64-bit value, whether the operand is int64/float64 or Inductor promoted the index dtype on its own. Those kernels are refused with a diagnosis instead of crashing the interpreter — see 64-bit Kernels on GCU. Removing the refusal (a compiler that gains 64-bit support) is what closes this, not a change on the torch_fl side
- TsingMicro is new here: the profile, the raw stream and the warp-size
correction are wired up and the graphs in
tests/integration/test_compile.pyare validated on a TX8, but nothing beyond them is. Two things are known and deliberately left alone: the TX runtime exposes no device-properties query, sotorch.flagos.get_device_propertiesreports constants matching FlagGems' own tsingmicro descriptor rather than values read from the hardware; and the Ascend-only Triton workarounds (triton_byte_loads.py,triton_resource_limits.py) are gated onis_cuda_like=False, so TsingMicro inherits them — they are unverified there. There is no CI lane for this vendor, so none of it runs automatically - TsingMicro's matmul is reduced-precision, and compiling changes which op
you get: a single
1x128 . 128x1dot on aflagostensor differs from a float32 CPU reference by ~1.6e-3, the signature of a TF32-shaped accumulation inside FlagGems'mmkernel. Eagernn.Lineardoes not show it because ATen sends it toaddmm, which is exact here; inductor instead schedulesextern_kernels.mmplus a bias add, so a compiled Linear is measurably further from the reference than the eager one (measured ~7e-4 relative on |out| ~ 2). Elementwise fusion stays bit-exact. This is the device's matmul precision rather than a compile defect — the samemmis what a hand-writtentorch.mmcalls — but it means compiled-vs-eager comparisons on this platform need a ~1e-3 relative tolerance, which is whattests/integration/test_compile.pyuses for it
- Phase 1: Inductor integration (flagos as a first-class GPU device)
- Phase 2: FlagTree integration — needs no shim; FlagTree substitutes for
triton at install time, so inductor picks it up unaided. Built from source
on H800 (
flagtree-0.6.0, Triton 3.6, nvidia backend). -
torch.compile(backend="flagos")end-to-end on FlagTree — forward and backward compile and match eager on NVIDIA and MetaX, with outputs and gradients remaining onflagos. The complete compile suite passes on both targets (the MetaX-specific event regression adds one case there). - Ascend via triton-ascend — experimental. Forward, backward, fused
elementwise and matmul+normalization compile and match eager on a real 910
(
Ascend910_9382, CANN 9.0.0, triton-ascend 3.2.0, torch 2.10.0+cpu);test_compile.pypasses 30/32 with a cold cache, the two skips being the FlagTree-only and MetaX-only cases. Superseded for the FlagGems route by FlagTree, and not re-measured there - FlagTree on Ascend, for the FlagGems eager route —
flagtree_ascend_policy.pygives FlagTree's Ascend backend a torch_npu-free policy andtorch_fl.flagosinstalls it during device init, so eagerly launched FlagGems kernels run on FlagTree0.6.2a1+ascend3.5; the route table was re-measured op by op on that stack (seedocs/reference/operator-support.md) - FlagTree on Ascend, for the
torch.compilepath — the policy is in place and the emitted C++ compiles against ATen with every strategy resolving through FlagTree's real registry, but the compile suite itself has not been run on FlagTree, so the three workarounds above are unverified there (#1046) - Benchmark fusion gains vs. stock inductor+triton on cuda
- Benchmark Ascend compile vs. the aclnn eager path
- Re-run the Ascend compile suite on FlagTree and retire the workarounds once its replacements land
- Phase 3: FlagGems-aware fusion (recognize pre-optimized patterns)
- Phase 4: Custom fusion patterns for flagos-specific ops
- PyTorch 2.0 torch.compile documentation
- TorchInductor overview
- FlagTree repository
- CPU torch + external libtorch_cuda.so — why several
torch.cudabindings need shimming torch_fl/compile/README.md— the registration surface in detail