feat(compile): torch.compile support for flagos as a first-class inductor GPU device

[lvyufeng]

Summary

Enables torch.compile on the flagos device by registering flagos with
TorchInductor as a first-class GPU device. The traced graph is handed to
compile_fx unchanged, and inductor emits Triton kernels that operate on flagos
tensors directly — no conversion to cuda, no copy at the graph boundary.

import torch, torch_fl
x = torch.randn(4096, 4096, device="flagos:0")
compiled = torch.compile(my_model, backend="flagos")
y = compiled(x)

This works because flagos runs on the physical GPU that torch.cuda describes:
its allocator delegates to c10::cuda::CUDACachingAllocator, so a flagos
tensor's storage already is CUDA memory.

Why not the device-aliasing approach

The first revision of this PR rewrote the graph and its example inputs to cuda
before calling compile_fx. That is not just a copy per call — it breaks
backward.

at::getAccelerator() is PrivateUse1/flagos in this build, and
torch::autograd::Node::stream() only yields a stream when a node's input device
type equals the accelerator. A cuda-rewritten graph therefore produces
stream-less autograd nodes, and AOT autograd's backward trace inside compile_fx
trips opt_ready_stream && opt_parent_stream (engine.cpp:1085). That was the
cause of 8 of the 11 test failures.

Confirmed by differential test: eager backward on plain cuda fails identically,
with torch.compile never involved.

Registration surface

What	Why
GPU_TYPES.append("flagos")	is_gpu() is a membership test on that list object; without it inductor takes the C++/CPU codegen path and never emits Triton.
Prime get_gpu_type()'s cache	It asserts at most one GPU type is available, and the torch.cuda shim reports available alongside flagos.
register_interface_for_device	Device state from torch.flagos, hardware properties from torch.cuda (same physical GPU).
DeviceProperties.create wrap	Reports flagos as cuda at the Triton boundary — Triton's NVIDIA backend hard-checks target.backend == "cuda", so a literal "flagos" finds 0 compatible backends. Inductor already does this in reverse for ROCm (hints.py:149).
register_device_op_overrides	Device guard / stream / synchronize snippets. Must directly inherit CUDADeviceOpOverrides — attributes present on the base class never reach __getattr__ delegation.
register_backend_for_device	The stock CUDA/Triton scheduling + wrapper codegen under the "flagos" key.

Two generated-kernel bugs that only surface under compilation

• detach re-dispatched into itself. The generated kernel called
  at::detach(self), which is also registered on PrivateUse1, so it dispatched
  straight back — infinite recursion. Eager hid this because DeviceBoxingGuard
  rewrites self's device metadata first; under FakeTensor it cannot, since the
  Python dispatch key sits above the backend key. Dynamo traces every
  nn.Linear through detach, so this was a stack-overflow segfault at trace
  time. Now emits at::native::detach (NATIVE_DIRECT_VIEW_OPS).
• gen_inplace didn't box optional<Tensor>. clamp_.Tensor handed unboxed
  flagos min/max to a CUDA self and crashed. Optionals are now materialized
  into holders, matching gen_functional_pure.

Both are covered by tests confirmed to fail (segfault at the exact asserting
line) against a build with the fixes reverted.

CPU-torch wheel accommodations

This build pairs a CPU-only pip torch with an external libtorch_cuda.so, so
several torch.cuda Python bindings are missing:

• CudaInterface.get_raw_stream re-attached — the binding exists, but the
  import-time _is_compiled() probe left it None.
• torch.cuda.memory_* routed to the flagos allocator backing the same pool.
• flagos Event/Stream in place of the dummy base classes.
• triton.cudagraphs = False (torch.cuda.CUDAGraph raises on construction) and
  use_static_cuda_launcher = False (not built).

flagos_compile_backend also accepts the mode/options/dynamic kwargs dynamo
forwards to named backends, expanding them into compile_fx config_patches
rather than mutating inductor's global config.

Testing

Rebased onto flagos/main (0700b61) and re-verified end to end on A100 in the
torch-fl-210 env:

Suite	Result
tests/integration/test_compile.py	12 passed, 1 skipped (was 2 passed / 8 failed)
tests/integration/ops/test_clamp_dispatch.py	15 passed
tests/integration/ops/ (full sweep)	489 passed, 160 skipped, 1 xfailed, 3 xpassed
tests/integration/ops/test_rng_dispatch.py	104 passed, 2 skipped, 1 xfailed
test_ops / allocator / factory / fallback_trace / clone_dispatch	135 passed
tests/unit	15 passed

The rebase brought in the RNG generator-injection work (#39, #49), which touches
the same scripts/codegen_ops.py templates; the conflict was resolved keeping
both, and python scripts/codegen_ops.py was verified to reproduce the committed
cuda_kernels.cc byte-for-byte.

Open work

• FlagTree integration to replace OpenAI Triton (scaffold present, gated by
  FLAGOS_USE_FLAGTREE=1, off by default)
• Fusion-gain benchmarking against stock inductor+triton on cuda — the previous
  revision of this PR quoted speedup figures that were measured under the old
  device-aliasing design, so they no longer describe this code and have been
  dropped rather than restated.
• Multi-GPU compilation not yet exercised

🤖 Generated with Claude Code