AI chips are flourishing, but PyTorch support across multiple chip platforms has long been trapped in siloed development. FlagOS Torch-FL introduces a virtual-device abstraction called flagos, reducing the work required to bring up a new chip to the implementation of only about 40 HAL functions. Operators can then be optimized incrementally, with the goal of making every type of AI accelerator as easy for developers to use as an NVIDIA GPU.
When a model moves from an NVIDIA GPU to another AI chip — an AI accelerator — developers face far more than a simple change of chip backend. Operator-fusion strategies, memory layouts, quantization precision, graph-compilation paths, and even the implementation of attention may all require architecture-specific development for different chips.
In other words, chip diversity is forcing what was once a general-purpose PyTorch inference framework into a collection of customized branches. Every new AI chip often requires the entire PyTorch integration layer to be rewritten: the runtime, memory management, event streams, and a vast number of operators must all be built from scratch. Three major obstacles now stand between different types of hardware and the PyTorch ecosystem:
-
Repeatedly reinventing the wheel
-
Conflicts between multiple in-process hardware plugins
-
High operator-maintenance costs This is precisely the part of cross-chip support for large models that consumes the most engineering effort and is the hardest to scale. To address this industry-wide challenge, FlagOS has launched Torch-FL, a unified PyTorch plugin for multi-chip support. Through a single virtual device,
flagos, it replaces the old siloed model — in which supporting one chip meant rewriting an entire integration stack — with a unified foundation where hardware can be brought up by implementing only about 40 HAL functions and operators can be optimized incrementally as needed. Torch-FL aims to end PyTorch fragmentation across diverse chip platforms.
Today, there are roughly three approaches to integrating different chips with PyTorch:
-
CUDA/HIP-compatible implementations (T-Head, MetaX, Hygon, and others): relatively low integration costs.
-
Non-CUDA-compatible implementations (Huawei Ascend, Cambricon, and others): siloed development, repeated implementation of shared modules, and enormous pressure to support operators.
-
Other hardware with custom ISAs (QingMicro, various XPUs, and others): likewise requires a complete, standalone integration plugin to be developed.
On the surface, the problem is the sheer amount of development work. Underneath are three structural conflicts that are difficult to avoid:
-
Massive duplication of effort: For foundational modules such as Runtime, Allocator, Guard, and Hooks, 95% of the logic implemented by different vendors is essentially the same, yet every vendor builds it independently.
-
PrivateUse1 exclusivity conflicts: PyTorch provides only one PrivateUse1 dispatch key. Plugins such as
torch_npuandtorch_musacannot coexist in the same process, making it inherently difficult to support real-world heterogeneous computing scenarios. -
Operator coverage is a bottomless pit: ATen has more than 2,000 operator schemas. The maintenance cost of handwritten kernels does not decrease over time: PyTorch version upgrades, new data types, and operator overloads all require continuous integration work.
The industry’s shared goal is simple: using different accelerator cards should offer the same PyTorch development experience as using NVIDIA GPUs. Existing mechanisms, however, make that goal difficult to achieve.
Torch-FL’s core innovation is to transform PrivateUse1 — which was originally private to a single vendor — into flagos, a general-purpose, reusable virtual-device layer with a three-tier architecture:
- Upper layer: one device name, `flagos`; one Python API; and one ProcessGroup.
-Middle layer: a configuration-driven per-op routing table rather than compile-time hard-coding.
- Lower layer: a pluggable vendor runtime plus three operator-provisioning paths.
The entire solution builds on the unified multi-chip foundation provided by FlagOS:
FlagGems: a high-performance, general-purpose operator library built with Triton/Triton-TLE.
FlagTree: an open-source, unified AI compiler based on Triton.
FlagCX: a high-performance communication component supporting same-chip, cross-chip, and supernode scenarios.
HAL: A Single Integration Boundary That Lowers the Entry Barrier by About 90%
Torch-FL consolidates hardware integration into the HAL abstraction layer, exposing a pure C ABI with only about 40 functions.
The minimum work required to support a completely new chip is to implement these 40 HAL functions. This generally requires only 3 files — device.cc, memory.cc, and stream.cc — and about 400 lines of code. Once that work is complete, the following can run normally:torch.randn(4,4,device=”flagos”)
The key design principle is that operators are not a prerequisite for hardware bring-up; they are targets for incremental optimization. When hardware is first integrated, computation uses CPU fallback. Functionality is correct, although performance is slower. Operator performance can then be improved iteratively as needed. Compared with the conventional plugin model, this fully decouples the runtime from operators and directly reduces the barrier to hardware integration by about 90%.
Three Operator Paths with Per-Op Mixed Dispatch in a Single Process
Simply bringing the hardware up is nowhere near enough. The ability to supply operators directly determines whether the hardware is usable in real workloads. Torch-FL does not bind operators to a single implementation path. Instead, it provides three operator-provisioning paths, allowing the backend to be selected freely for each operator within the same process. If an operator bug appears, a one-line configuration change can switch it back to an alternative implementation without recompiling the project.
1. FlagGems Triton: One kernel can run across multiple chip platforms. As long as a hardware vendor provides a Triton backend, the operator library can be reused directly. Two integration modes are available: the Python layer offers broader coverage, while the C++ runtime avoids the GIL and delivers lower latency. It currently covers 18 high-demand operators.
2. CUDA Boxing: Zero-copy metadata rewriting allows a flagos tensor and a CUDA tensor to share the same physical GPU memory. Only the device metadata inside TensorImpl is modified, while Guard provides exception safety. All 18,679 lines of routing code are generated automatically by Codegen, with 0 handwritten lines, covering 2,036 operator routes. This path is suitable for CUDA/HIP-compatible hardware.
3. Native Codegen: A fallback for hardware that cannot reuse CUDA. For hardware such as Ascend, GCU, and MUSA, where the CUDA ecosystem is unavailable, this path connects directly to vendors’ native operator libraries. Codegen is a purpose-built generator with symbol validation — not a simple template wrapper — ensuring correct function calls.
A Lightweight In-House Dispatcher with 93% of Its Code Generated by Codegen
Multi-path operator dispatch places new demands on the framework’s dispatch layer. The set of backends supported by PyTorch’s native DispatchStub is fixed at compile time, so it cannot dynamically switch among multiple backends at runtime. Torch-FL therefore replaces it with a lightweight in-house Dispatcher:
Engineering statistics show that Torch-FL contains 48,124 lines of C++ code in total. Codegen automatically generates 93% of it, leaving only about 3,200 lines of handwritten C++. This is a key reason the project can rapidly expand to eight hardware platforms.
Deeply Engineered Subsystems That Solve Real-World Problems
Even after the dispatch and operator frameworks are in place, numerous practical engineering pitfalls remain. Many integration solutions can run a demo, only to expose hidden bugs when deployed for training or multi-card distributed workloads. Torch-FL provides comprehensive engineering fixes and support for easily overlooked low-level subsystems such as memory management, random-number generation, and distributed communication:
Allocator: Resolves inaccurate GPU-memory accounting and double reservation caused by dual memory pools in Boxing mode, while remaining compatible with header and symbol conflicts in DCU’s HIP implementation.
RNG: Fixes the architectural flaw in which torch.manual_seed affects only some operators, providing unified random-number behavior across the entire execution path.
Distributed ProcessGroupFlagOS: Implements a native ProcessGroup subclass without MonkeyPatch. Communication backends are prioritized as FlagCX → HCCL → NCCL, with zero-copy views delegating to the underlying communication capabilities. This follows the principle of “implementing interfaces at framework extension points rather than modifying upper-layer APIs.”
Deep Integration with torch.compile: A Two-Phase Roadmap
Once hardware and operator support for Eager mode is complete, the next step is to integrate with PyTorch’s compilation capabilities and fully unlock the hardware’s compute performance. Torch-FL adopts a phased approach to supporting torch.compile:
Phase 1: Register the flagos backend with torch._dynamo. TorchInductor generates fused Triton/CUDA kernels while reusing Torch-FL’s complete dispatch path.
Phase 2: Replace OpenAI Triton with FlagTree as the multi-backend compiler, allowing the same Triton kernel to be compiled and run on many types of hardware, including NVIDIA, Huawei Ascend, Cambricon, and MetaX.
At compile time, a flagos tensor is presented as a CUDA tensor; at runtime, zero-copy metadata conversion is performed. Tests of an MLP model on A100 hardware show that native CUDA Inductor delivers a 2.8× speedup, while flagos Inductor achieves a 2.6× speedup — 93% performance parity.
Beyond matching standard PyTorch capabilities, Torch-FL provides 2 distinctive capabilities that are difficult to implement directly in native PyTorch, covering edge hardware and low-precision large-model workloads:
1. Support for embodied-AI chips: In Eager execution mode, CPU fallback ensures functional availability. With full-graph compilation through torch.compile, the model is offloaded to the BPU for hardware acceleration, expanding the use of PyTorch on edge devices for embodied intelligence.
2. FP8/FP4 software emulation: For chips whose hardware does not natively support low precision, Torch-FL provides software emulation and can natively load FP8 and FP4 pretrained large-model weights from NVIDIA.
After a series of low-level design improvements and extensive engineering work, Torch-FL has completed verification across multiple types of hardware, and its production-readiness has been validated in real tests.
On September 20, Alibaba open-sourced Qwen-Image-2.1. Using the complete Torch-FL technology stack, the FlagOS community completed cross-chip support and inference validation in a Day 0 release for Qwen-Image-2.1. The work covered 8 chip platforms: T-Head, NVIDIA, Moore Threads, Huawei Ascend, MetaX, Hygon, Enflame, and Arm. Inference accuracy on every platform matched the native NVIDIA version. Standard Diffusers inference code ran with zero code changes, and on multiple platforms, general-purpose operators delivered out-of-the-box performance that exceeded vendors’ existing implementations. FlagOS multi-chip support is now included in the Hardware Support section of the official Qwen-Image-2.1 GitHub README.
Unlike language models, open-source generative multimodal models are released natively in the form of PyTorch + Diffusers. To make them available across chip platforms on the day of release, support must be implemented directly at the PyTorch layer. Without changing the native PyTorch or Diffusers interfaces, Torch-FL extends model inference from NVIDIA platforms to multiple AI chips. Of the 62 operators in the NVIDIA inference path, 37 were handled by FlagGems. The same plugin and operator path was reused across 8 chip platforms. On several platforms where vendors had not yet performed model-specific optimization, out-of-the-box performance surpassed the vendors’ existing PyTorch paths. Accuracy was evaluated using the T2I-100 and COCO text-to-image benchmarks; on every platform, the difference in CLIPScore from native NVIDIA execution was within ±5%.
Torch-FL has now completed integration and validation with 7 cloud-chip vendors and 1 edge embodied-AI chip vendor. It can run more than 40 mainstream large models, using the same test code to validate operator correctness, Qwen3 training and inference, and multi-card DDP.
The project is open-sourced under the Apache-2.0 license. The community is welcome to participate in testing, report issues, and submit pull requests.
Open-source repository: https://github.com/flagos-ai/Torch-FL The AI hardware ecosystem is not just a competition in chip performance. Framework support and the software foundation are equally critical.
Torch-FL aims to break away from the siloed model of “one PyTorch plugin per chip.” Through the flagos virtual device, a minimal HAL integration layer, and multi-path operator dispatch, it lowers the overall cost of bringing different types of hardware into PyTorch — enabling more developers to train and run inference for AI models smoothly across a wide range of accelerator hardware.
As hardware innovation flourishes, the software foundation must lead the way. Try FlagOS Torch-FL and join us in building a unified PyTorch ecosystem for diverse hardware.
Must-Know Tips| Ending PyTorch Fragmentation Across Diverse Chips: How FlagOS Torch-FL Makes AI… was originally published in Stackademic on Medium, where people are continuing the conversation by highlighting and responding to this story.