From b4eb99f869ce8ddb80cf77164d337a989cca3f5d Mon Sep 17 00:00:00 2001 From: "nate.river" Date: Thu, 6 Aug 2026 11:49:38 +0800 Subject: [PATCH] feat(gcu): run FlagGems Triton kernels on Enflame GCU Enflame GCU has no C++ FlagGems path (there is no liboperators.so for the tops stack, so FLAGGEMS_KERNEL/FLAGGEMS_PYTHON stay off), which left it on topsaten kernels and cpu_fallback only. This registers FlagGems' Triton kernels through the Python layer instead, calling flag_gems.enable() straight onto PrivateUse1: 666 aten ops, 244 of them the vendor's own gcu300 overrides. The kernels reach the device through Enflame's triton_gcu plugin, which was written against the vendor's torch_gcu plugin and names PrivateUse1 "gcu". torch_fl claims PrivateUse1 first as "flagos", and a process can rename it only once, so torch_fl/accelerator/gcu/_gcu_compat.py redirects the vendor backend: the device name in its toolkit/backend/driver, a wall-clock timing Event for do_bench (flagos.Event derives from torch.cuda.Event, a dummy on a CPU-only wheel), and the two hardcoded device-0 lookups that made a kernel on flagos:1 launch against device 0. Two mismatches need correcting on the FlagGems side because an op silently calls something other than what it looks like: * bind_vendor_ops_in_generic_modules -- generic ops import sub-ops by value at module scope (linear_backward does `from .mm import mm`), so the binding survives the vendor override. The generic mm widens its indices to int64, which the GCU backend refuses to legalize, so every backward through nn.Linear failed to compile while torch.mm worked. * device_guarded_config -- FlagGems allocates intermediates on input.device but launches Triton on the current device. This is the c10::OptionalDeviceGuard that python_op_caller.cc applies, in Python, at the one place these ops go through. It also diverts all-int64 arithmetic to the CPU kernel, since the tops stack has no int64. _GCU_EXCLUDED_OPS lists the ops whose FlagGems kernel the GCU cannot compile or compiles wrongly, each verified individually on hardware and each left to reach topsaten or cpu_fallback; the comments record what the defect is. The conv family is excluded even though every forward is correct, because its VJP trips a SIP assert that aborts the process instead of raising. _flaggems_exclusion_names() translates aten names into what the exclusion filter actually matches -- flag_gems.enable(unused=...) compares the *implementing function* name, so "normal.Tensor_float" excludes nothing (the function is normal_tensor_float) and the op gets registered anyway. The whole path is optional: when triton_gcu or its /opt/triton_gcu toolchain is absent, registration is skipped and the build behaves exactly as before. A registration failure warns rather than breaking `import torch_fl`. Test suite is at parity with FlagGems on and off: 521 passed, 22 failed, 9 errors (conv/RNG/inductor/transformers, all pre-existing and identical in both runs). Verified on gcu300, single- and multi-device. --- setup.py | 12 + torch_fl/__init__.py | 240 ++++++++++- torch_fl/accelerator/gcu/__init__.py | 13 + torch_fl/accelerator/gcu/_gcu_compat.py | 545 ++++++++++++++++++++++++ torch_fl/flagos/__init__.py | 52 ++- 5 files changed, 859 insertions(+), 3 deletions(-) create mode 100644 torch_fl/accelerator/gcu/__init__.py create mode 100644 torch_fl/accelerator/gcu/_gcu_compat.py diff --git a/setup.py b/setup.py index 8c468a41..de46c287 100644 --- a/setup.py +++ b/setup.py @@ -383,10 +383,22 @@ def build_deps(): # layer and libtopsaten the operators, so every CUDA/vendor kernel set # stays off and GCU_KERNEL (topsaten) provides the compute ops. Ops # without a topsaten kernel fall back to CPU. + # + # FLAGGEMS_PYTHON is off because the C++ dispatch path it builds + # (python_op_caller.cc + generated/flaggems_python_kernels.cc) routes + # per-op through backends_*.conf, and gcu has its own conf naming the + # topsaten kernels. FlagGems still reaches the GCU here -- just through + # the Python layer instead, where torch_fl calls flag_gems.enable() + # directly onto PrivateUse1 (see torch_fl/accelerator/gcu/_gcu_compat.py + # and _register_flaggems_operators). That needs Enflame's triton_gcu + # plugin plus the /opt/triton_gcu toolchain; when either is missing the + # registration is skipped and everything stays on topsaten, so this + # build works with or without them. cmake_args.extend( [ "-DCUDA_KERNEL=OFF", "-DFLAGGEMS_KERNEL=OFF", + "-DFLAGGEMS_PYTHON=OFF", "-DMETAX_KERNEL=OFF", "-DASCEND_KERNEL=OFF", "-DGCU_KERNEL=ON", diff --git a/torch_fl/__init__.py b/torch_fl/__init__.py index 62d625da..8adb7445 100644 --- a/torch_fl/__init__.py +++ b/torch_fl/__init__.py @@ -524,6 +524,21 @@ def _patch_flaggems_codegen_config(): os.environ.setdefault("GEMS_VENDOR", "hygon") return + # --- Enflame GCU branch --- + # Keyed on the build accelerator for the same reason as DCU: no runtime probe + # distinguishes GCU here, and the tops stack has no libcuda.so, so without + # this branch GCU would reach the ascend fallback and get GEMS_VENDOR=ascend + # (which also picks the wrong comm profile). FlagGems' Triton kernels need + # Enflame's triton_gcu plugin plus its /opt/triton_gcu compiler toolchain; if + # either is missing, patch_triton_gcu_for_flagos() returns False and we leave + # GEMS_VENDOR unset so the topsaten kernels and cpu_fallback stay in charge. + if _build_accelerator() == "gcu" and os.environ.get("GEMS_VENDOR") != "ascend": + from torch_fl.accelerator.gcu._gcu_compat import patch_triton_gcu_for_flagos + + if patch_triton_gcu_for_flagos(): + os.environ.setdefault("GEMS_VENDOR", "enflame") + return + # --- Generic NVIDIA CUDA branch (default) --- if ( os.environ.get("FLAGOS_DISABLE_CUDA_SHIM", "0") != "1" @@ -648,13 +663,47 @@ def _patched_cuda_device_init(self, device): "randperm", "empty.memory_format", # Already registered in C++ "empty_strided", # Already registered in C++ - # Random ops that use device context + # FlagGems registers the *bare* aten name "empty" as well, implemented by a + # function also called "empty" -- so "empty.memory_format" above does not + # filter it (the registrar matches on the function name). Left registered, it + # is the one factory op FlagGems takes over, and it ignores the requested + # device: an `empty(..., device="flagos:1")` while flagos:0 is current + # allocates on 0, and the first write through that pointer faults the driver + # ("IoctlCmdWriteRead errno[14]. The device is out of service", then a SIP + # exception that aborts the process). The C++ empty is already correct here. + "empty", + # Random ops that use device context. + # + # These also cover FlagGems' philox path, which reads the generator state as + # exactly two int64s (the CUDA layout: seed + offset). torch_fl's flagos + # generators are CPU Mersenne-Twister generators whose state unpacks to 632 + # int64s, so philox_backend_seed_offset raises "too many values to unpack". + # Every name FlagGems registers for an RNG op has to appear here or the op + # reaches that unpacking; the factory/RNG ops are correct on the topsaten and + # CPU paths anyway, so nothing is lost by keeping them off FlagGems. "uniform_", + "normal_", "normal.float_Tensor", "normal.Tensor_float", "normal.Tensor_tensor", + "normal.Tensor_Tensor", + "log_normal_", + "randint", + "randint_like", "exponential_", "multinomial", + # The rest of the philox users, found by running tests/integration/ops/ + # test_rng_dispatch.py: each of these reaches philox_backend_seed_offset and + # fails the same way. cauchy_/dropout/poisson are the *vendor's* gcu300 + # overrides, so they are not discoverable from flag_gems/ops/ alone. + "bernoulli", + "bernoulli.p", + "bernoulli_.float", + "bernoulli_.Tensor", + "cauchy_", + "dropout", + "native_dropout", + "poisson", # Copy ops - already registered in C++, skip to avoid duplicate registration "copy_", "_to_copy", @@ -702,6 +751,109 @@ def _patched_cuda_device_init(self, device): } +# Ops excluded from FlagGems on Enflame GCU only, on top of _EXCLUDED_OPS. +# +# These have a FlagGems kernel that its Triton backend cannot compile for the +# GCU, so they must stay unregistered to keep reaching the topsaten kernels or +# cpu_fallback. Verified individually on hardware -- only the listed overload +# fails, e.g. var.dim, std.correction and var_mean.correction all work. +_GCU_EXCLUDED_OPS = { + # NOTE: FlagGems matches this list against the *implementing function* name + # (op_registrar.config_filter compares item[1].__name__), not the aten op + # name -- those merely coincide for most ops. aten::var.correction is + # implemented by var_correction, so that is the name to list. + # + # The full-reduction path (var_kernel_1, var.py:88) emits an int64 widening + # that the GCU backend marks illegal: "failed to legalize operation + # 'arith.extsi'" -- consistent with the tops stack having no int64 kernels. + "var", + "var_correction", + "var_dim", + "var_mean", + # Both are built on Triton's float `%`, which on the GCU returns x rather + # than 0 when y divides x exactly (the internal division lands just below + # the integer, so the floor is one too low). torch.remainder(2*y, y) then + # gives y instead of 0 -- for ~10% of random lanes, silently. Non-multiple + # operands are correct, which is why this needs an exact-multiple probe to + # see. Verified on gcu300 with the vendor rem_tt/fmod kernels. + "remainder", + "remainder_", + "fmod_scalar", + "fmod_tensor", + "fmod_scalar_", + "fmod_tensor_", + "fmod_", + # Same rounding defect seen through the quotient instead of the remainder: + # floor_divide(y, y) yields 0 rather than 1 on those same lanes. + "floor_divide", + "floor_divide_", + # No GCU kernel to link against: the vendor linker rejects the relocation + # for tops_nv_nextafterf_v4_fp32 ("R_DTU_ADDR16_LO_ICALL cannot be used + # against symbol"), so the op cannot be compiled at all. + "nextafter", + "nextafter_", + # The sort kernels emit the same illegal int64 widening as var_kernel_1 + # ("failed to legalize operation 'arith.extsi'"). msort is sort's caller, so + # it fails identically, and sort.stable (function sort_stable) shares the + # kernel -- it fails as "Pipeline run failed: PassManager execution failed". + # topk/argmax use different kernels and are fine. + "sort", + "sort_stable", + "msort", + # stack.py:65 builds its offsets in int64 and hits that same legalization + # failure -- and the backend then core-dumps while reporting the error, so + # this one cannot be left to raise. cat/vstack/hstack are unaffected and + # verified correct. + "stack", + # The conv VJP (flag_gems/ops/conv2d.py, which conv1d unsqueezes into) + # passes a stride of 0 as a runtime argument, and the GCU asserts on that + # inside the kernel: "Not Support dynamic stride is 0, please add + # tl.constexpr to stride arg in kernel args list". That is a SIP assert, so + # it aborts the process instead of raising -- it cannot be caught and fallen + # back from, which is why the whole family stays off FlagGems even though + # every forward is numerically correct. + "conv1d", + "conv2d", + "conv3d", + "conv_transpose1d", + "conv_transpose2d", + "_conv_depthwise2d", + "cudnn_convolution", + # The vendor's own gcu300 layernorm.py:326 hits the int64 widening in its + # backward kernel, which breaks any .backward() through a LayerNorm. The + # forward (layer_norm) compiles and is verified correct, so only the backward + # is excluded -- the gradient falls back to the topsaten/CPU path. + "native_layer_norm_backward", + "layer_norm_backward", + # int64 again, from the other direction: these two are asked for int64 + # *output* rather than int64 indices. `zeros(dtype=torch.int64)` goes through + # the vendor's gcu300 zeros.py zero_ and fails to compile; `scalar_tensor` + # with dtype=int64 compiles but returns garbage (42 came back as + # 4441830098096545884). Both are correct on the C++/topsaten path. + "zero_", + "scalar_tensor", + # flag_gems/ops/diff.py builds its offsets in int64 -- same legalization + # failure as stack. torch.diff decomposes to narrow+sub on the fallback path. + "diff", + # Same story as layernorm: the vendor's own gcu300 embedding.py:98 backward + # kernel does not legalize, while its forward (embedding) is correct and + # stays on FlagGems. + "embedding_dense_backward", + "embedding_backward", + # fill_ silently corrupts int64 tensors -- fill_(42) on an int64 tensor comes + # back as -4846589848703729622, at every rank, with no error. The int64 + # diversion in device_guarded_config cannot rescue it: fill_ writes through + # its operand, so computing on a CPU copy would discard the result. Excluding + # it also fixes torch.scalar_tensor(dtype=torch.int64), which is an + # empty+fill_ underneath and returned garbage even though scalar_tensor + # itself was already excluded. + "fill_.Scalar", + "fill_.Tensor", + "fill.Scalar", + "fill.Tensor", +} + + # Cache for CUDA runtime library _cudart_lib = None _cudaMemcpy = None @@ -734,6 +886,38 @@ def _get_cudaMemcpy(): return _cudaMemcpy +def _flaggems_exclusion_names(flag_gems, aten_names): + """Translate aten op names into the function names FlagGems excludes on. + + ``flag_gems.enable(unused=...)`` looks like it takes aten op names, but its + registrar filters on the *implementing function* name + (``op_registrar.config_filter`` compares ``item[1].__name__``). Those + coincide for plain ops -- "randn", "mm" -- and diverge for overloads and + private ops: ``normal.Tensor_float`` is implemented by + ``normal_tensor_float``, ``_softmax`` by ``softmax``, ``div.Scalar`` by + ``true_divide``. Passing an aten name that diverges silently excludes + nothing, and the op gets registered after all. + + Both spellings are returned: the function name is what actually filters, + while keeping the original is harmless and covers ops whose two names agree. + Names absent from FlagGems' config pass through unchanged. + """ + op_to_func = {} + for item in getattr(flag_gems, "_FULL_CONFIG", ()): + if len(item) < 2: + continue + func_name = getattr(item[1], "__name__", None) + if func_name: + op_to_func.setdefault(item[0], func_name) + + names = set(aten_names) + for aten_name in aten_names: + func_name = op_to_func.get(aten_name) + if func_name: + names.add(func_name) + return sorted(names) + + def _register_flaggems_operators(): """ Register FlagGems operators with the PrivateUse1 (flagos) dispatch key. @@ -753,6 +937,60 @@ def _register_flaggems_operators(): # flag_gems not installed, will use cpu_fallback return 0 + # Enflame GCU has no C++ FlagGems path (FLAGGEMS_KERNEL is off -- there is no + # liboperators.so for the tops stack), so the Python registration is the only + # way its Triton kernels get used. Everything not registered here stays on + # the topsaten kernels or reaches cpu_fallback, as before. + if _build_accelerator() == "gcu": + from torch_fl.accelerator.gcu._gcu_compat import is_triton_gcu_available + + if not is_triton_gcu_available(): + return 0 + try: + import flag_gems + + from torch_fl.accelerator.gcu._gcu_compat import ( + bind_vendor_ops_in_generic_modules, + device_guarded_config, + patch_flaggems_device_name, + ) + + patch_flaggems_device_name() + bind_vendor_ops_in_generic_modules(flag_gems) + excluded = _flaggems_exclusion_names( + flag_gems, _EXCLUDED_OPS | _GCU_EXCLUDED_OPS + ) + _flaggems_lib = torch.library.Library("aten", "IMPL") + # enable() reads _FULL_CONFIG off the module, so the guarded table + # is swapped in for the call and restored right after -- anything + # else reading _FULL_CONFIG later sees the unwrapped functions. + original_config = flag_gems._FULL_CONFIG + flag_gems._FULL_CONFIG = device_guarded_config(flag_gems) + try: + flag_gems.enable(lib=_flaggems_lib, unused=excluded) + finally: + flag_gems._FULL_CONFIG = original_config + # What FlagGems actually took over, i.e. the same filter enable() + # applied: the aten names whose implementing function survived it. + skip = set(excluded) + _registered_ops = sorted( + entry[0] + for entry in flag_gems._FULL_CONFIG + if getattr(entry[1], "__name__", None) not in skip + ) + return 1 + except Exception as exc: + # A broken vendor Triton must not take down `import torch_fl`: the + # topsaten kernels and cpu_fallback are a complete, correct path. + import warnings + + warnings.warn( + f"FlagGems registration for GCU failed ({type(exc).__name__}: " + f"{exc}); falling back to the topsaten kernels.", + stacklevel=2, + ) + return 0 + _flaggems_lib = torch.library.Library("aten", "IMPL") _registered_ops = [] return 0 diff --git a/torch_fl/accelerator/gcu/__init__.py b/torch_fl/accelerator/gcu/__init__.py new file mode 100644 index 00000000..a491ed27 --- /dev/null +++ b/torch_fl/accelerator/gcu/__init__.py @@ -0,0 +1,13 @@ +# 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. diff --git a/torch_fl/accelerator/gcu/_gcu_compat.py b/torch_fl/accelerator/gcu/_gcu_compat.py new file mode 100644 index 00000000..1f0d0be9 --- /dev/null +++ b/torch_fl/accelerator/gcu/_gcu_compat.py @@ -0,0 +1,545 @@ +# 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. + +""" +FlagGems compatibility layer for Enflame GCU. + +FlagGems' Triton kernels reach the GCU through Enflame's ``triton_gcu`` plugin +(the vendor Triton backend; ``pip install triton-gcu`` from the FlagOS enflame +index, plus the ``triton-gcu`` deb that installs the ``/opt/triton_gcu`` +compiler toolchain). That plugin was written against Enflame's own ``torch_gcu`` +plugin, which claims PrivateUse1 and names the device ``gcu``. torch_fl claims +PrivateUse1 first and names it ``flagos``, and a process can rename PrivateUse1 +only once -- so every place triton_gcu says "gcu" has to be redirected here. + +None of this changes what the kernels compute; it only makes the vendor backend +agree with torch_fl about what the device is called and how to time a kernel. + +Requires, on top of the shims below: + * ``GEMS_VENDOR=enflame`` -- FlagGems' vendor autodetect keys off + ``hasattr(torch, "gcu")``, which torch_fl does not provide. Set + automatically once the vendor stack is confirmed present. + * ``COMPILE_ARCH=gcu300``/``gcu400`` -- selects the arch without asking the + driver for a stream handle (see GCUDriver.__init__). Set automatically from + the driver's reported arch when unset. + * a libstdc++ new enough for FlagGems' sqlite3/sqlalchemy import -- conda's + ``$CONDA_PREFIX/lib/libstdc++.so.6`` works where a system one may not. + +Two further mismatches are corrected here rather than in the vendor stack, both +because a FlagGems op silently calls something other than what it looks like: +``bind_vendor_ops_in_generic_modules`` (a generic op reaching a generic sub-op +past the vendor's override) and ``device_guarded_config`` (a kernel launching on +the current device rather than its operand's). +""" + +import functools +import os +import time + + +def is_triton_gcu_available() -> bool: + """True when both halves of the vendor Triton stack are installed. + + The Python plugin alone is not enough: kernel compilation shells out to + ``$TRITON_GCU_PATH/bin/gcu-compiler-{opt,compile}``, which ship in a + separate deb. Checking both here keeps the failure at import time (where it + is actionable) instead of at the first kernel launch. + """ + import importlib.util + + if importlib.util.find_spec("triton_gcu") is None: + return False + if importlib.util.find_spec("triton") is None: + return False + datadir = os.environ.get("TRITON_GCU_PATH") or "/opt/triton_gcu" + return all( + os.path.exists(os.path.join(datadir, "bin", tool)) + for tool in ("gcu-compiler-opt", "gcu-compiler-compile") + ) + + +class _WallClockEvent: + """Wall-clock stand-in for a CUDA-style timing Event. + + ``triton.testing.do_bench`` -- which FlagGems' autotuner calls -- needs + ``Event(enable_timing=True)`` with ``record``/``elapsed_time``. torch_fl's + ``flagos.Event`` derives from ``torch.cuda.Event``, which on a CPU-only + torch wheel is a dummy base class that raises on instantiation. + + Autotuning only compares candidate configs against each other, so relative + wall-clock timings pick the same winner. Every GCU op path synchronizes + before returning, so the measured interval does bracket real device work. + """ + + def __init__(self, enable_timing: bool = True): + self._t = None + + def record(self, *args, **kwargs): + self._t = time.perf_counter() + + def synchronize(self): + pass + + def wait(self, *args, **kwargs): + pass + + def query(self) -> bool: + return True + + def elapsed_time(self, end) -> float: + if self._t is None or end._t is None: + return 0.0 + return (end._t - self._t) * 1000.0 + + +def _patch_vendor_device_name() -> None: + """Make FlagGems' enflame descriptor report the device as ``flagos``. + + FlagGems' _enflame backend declares ``device_name="gcu"``, which surfaces as + ``flag_gems.runtime.device.name``. Its ops use that string two ways -- as a + torch device for intermediate allocations (``torch.empty(..., device=device)``) + and as an identity check against their inputs + (``assert X.device.type == device``, in 20 op files including maximum/minimum + and the upsample family). Both want the name torch_fl actually registered. + + Patched on ``VendorDescriptor`` before ``flag_gems`` is imported, because + ``DeviceDetector`` copies the value out of the descriptor at construction and + is a singleton -- after that first import the name is fixed. ``backend_utils`` + is a top-level module, so importing it does not pull in flag_gems itself. + """ + try: + import backend_utils + except ImportError: + return + + descriptor = getattr(backend_utils, "VendorDescriptor", None) + if descriptor is None: + # flag_gems < 5.3 calls it VendorInfoBase. + descriptor = getattr(backend_utils, "VendorInfoBase", None) + if descriptor is None or getattr(descriptor, "_flagos_patched", False): + return + + original_init = descriptor.__init__ + + def _init(self, *args, **kwargs): + if kwargs.get("device_name") == "gcu": + kwargs["device_name"] = "flagos" + original_init(self, *args, **kwargs) + + descriptor.__init__ = _init + descriptor._flagos_patched = True + + +def patch_triton_gcu_for_flagos() -> bool: + """Redirect Enflame's triton_gcu backend onto torch_fl's flagos device. + + Returns False (having changed nothing) when the vendor stack is not + installed, so callers can fall through to the topsaten/CPU paths. + """ + if not is_triton_gcu_available(): + return False + + import torch + + from torch_fl import flagos + + # 1. triton_gcu's driver calls torch.gcu.current_device()/current_stream(). + # torch_fl registers the same surface as torch.flagos. + if not hasattr(torch, "gcu"): + torch.gcu = flagos + + # 2. toolkit/backend build torch device strings from a module-level + # device_name = "gcu". PrivateUse1 is already named flagos. + import triton_gcu.triton.backend as _backend + import triton_gcu.triton.toolkit as _toolkit + + _toolkit.device_name = "flagos" + _backend.device_name = "flagos" + + # 3. The autotuner's L2-flush buffer is allocated with a hardcoded + # device='gcu' (driver.py get_empty_cache_for_benchmark). + # + # The index matters as much as the name. ``do_bench`` allocates this + # buffer once and then calls ``clear_cache(cache)`` -- i.e. ``cache + # .zero_()`` -- between timing runs, interleaved with the kernel it is + # autotuning. Written as device="flagos" the buffer lands on whatever + # device was current at allocation, so autotuning an op on flagos:1 + # while flagos:0 is current has zero_ writing device-0 memory from a + # device-1 context. The tops runtime does not reject that: the SIP + # faults asynchronously and the process aborts at the next + # synchronization ("Receive Sip error message" / "Detected context + # error!!!"), several ops after the one that caused it. Naming the + # current device pins the buffer to the same device as the kernel. + import triton_gcu.triton.driver as _driver + + _driver._GCUDriver.get_empty_cache_for_benchmark = lambda self: torch.empty( + 256, dtype=torch.int, device=torch.device("flagos", flagos.current_device()) + ) + + # 4. Give do_bench a usable timing Event (see _WallClockEvent). + flagos.Event = _WallClockEvent + torch.gcu = flagos + + # 5. Setting COMPILE_ARCH (below) takes a GCUDriver constructor branch that + # hardcodes `get_current_device = lambda: 0`, so a kernel operating on + # flagos:1 would launch against device 0 and read another device's memory + # -- in practice it hangs, since a tops pointer only resolves against the + # current device (see the tops-pointers-are-device-scoped note). + # + # Both driver classes need fixing: _GCUDriver.__init__ copies the inner + # GCUDriver's bound method onto the instance, so patching only the wrapper + # class would be undone by the next construction. get_current_stream keeps + # the branch's 0 (the default stream), which is what the launcher wants. + # get_active_torch_device also builds a torch.device("gcu", ...). + # Patching the classes is not enough, and this is the subtle part: + # GCUDriver.__init__ assigns `self.get_current_device = lambda idx=0: 0` + # as an *instance attribute*, which shadows anything set on the class, and + # _GCUDriver.__init__ then copies that attribute onto itself. _GCUDriver + # caches only its instance (in __new__), so __init__ -- and with it a fresh + # GCUDriver -- re-runs on every `_GCUDriver()` call, undoing a class patch. + # So wrap both constructors and re-assign the attribute afterwards. + # + # Symptom when this is wrong: >=3 kernel launches on device 0 followed by + # one on another device kills the driver outright ("Receive Sip error + # message", then "Receive Abort message from KMD: Sip exception"), taking + # the process with it. Fewer than three launches survive, which is what + # made this look like a bug in whichever op happened to run first. + _current_device = staticmethod(lambda: flagos.current_device()) + _driver.GCUDriver.get_current_device = _current_device + _driver._GCUDriver.get_current_device = _current_device + _driver._GCUDriver.get_active_torch_device = lambda self: torch.device( + "flagos", flagos.current_device() + ) + + def _patch_init(cls): + original = cls.__init__ + if getattr(original, "_flagos_patched", False): + return + + @functools.wraps(original) + def __init__(self, *args, **kwargs): + original(self, *args, **kwargs) + # Drop the instance attribute so the class-level staticmethod above + # is what lookups find. get_current_stream keeps the constructor's 0 + # (the default stream), which is what the launcher wants. + self.__dict__.pop("get_current_device", None) + + __init__._flagos_patched = True + cls.__init__ = __init__ + + _patch_init(_backend.GCUDriver) + _patch_init(_driver._GCUDriver) + + # An instance built before this point still holds the copied attribute. + _existing = getattr(_driver._GCUDriver, "instance", None) + if _existing is not None: + _existing.__dict__.pop("get_current_device", None) + if getattr(_existing, "_driver", None) is not None: + _existing._driver.__dict__.pop("get_current_device", None) + + # COMPILE_ARCH lets GCUDriver skip the torch.gcu stream lookup during its + # own construction. Derive it from the arch the driver reports + # ("dtu-enflame-tops--gcu300" -> "gcu300") rather than hardcoding a chip. + if "COMPILE_ARCH" not in os.environ: + import re + + try: + arch = _driver._GCUDriver().get_arch() + match = re.search(r"gcu\d+", arch) + if match: + os.environ["COMPILE_ARCH"] = match.group(0) + except Exception: + # Leave COMPILE_ARCH unset: the driver then resolves the arch from + # the live device, which works once the shims above are in place. + pass + + return True + + +def patch_flaggems_device_name() -> bool: + """Point FlagGems' own device name at ``flagos``. + + FlagGems takes the device name from its enflame vendor descriptor, so + ``flag_gems.runtime.device.name`` is ``"gcu"`` while tensors here report + ``"flagos"``. Ops that compare the two (``maximum``/``minimum`` do + ``assert X.device.type == device``) then fail on correct input. + + Unlike the triton_gcu shims this runs *after* FlagGems is imported, so it + has two things to fix. ``DeviceDetector`` is a singleton, so correcting + ``.name`` on it covers every later reader. But the op modules do + ``device = device.name`` at module scope, and importing any part of + ``flag_gems.runtime`` eagerly imports them -- those already hold the old + literal, so each such module global is rewritten as well. + + Call this after ``import flag_gems`` and before ``flag_gems.enable()``. + + Returns False if FlagGems is not installed or is not on the enflame vendor. + """ + import importlib.util + import sys + + if importlib.util.find_spec("flag_gems") is None: + return False + try: + from flag_gems.runtime.backend.device_finder import DeviceDetector + except ImportError: + # Older FlagGems releases keep DeviceDetector in backend.device. + try: + from flag_gems.runtime.backend.device import DeviceDetector + except ImportError: + return False + + detector = DeviceDetector() + if detector.vendor_name != "enflame": + return False + stale = detector.name + if stale == "flagos": + return True + detector.name = "flagos" + # Not filtered by module name: the vendor's arch-specific overrides are + # loaded under bare keys like "gcu300.ops.maximum", outside the flag_gems + # package path. Match on the stale value instead, which is specific enough. + for module in list(sys.modules.values()): + if module is not None and getattr(module, "device", None) == stale: + module.device = "flagos" + return True + + +def _is_int64_arithmetic(args, kwargs) -> bool: + """True when every tensor operand is int64, i.e. int64 *arithmetic*. + + The GCU has no int64 kernels, which FlagGems' own enflame descriptor states + (``int64_enabled=False``) and then never enforces -- nothing reads + ``DeviceDetector.support_int64``. An int64 tensor therefore reaches a Triton + kernel that cannot legalize the type: a compile error at best ("failed to + legalize operation 'arith.extsi'"), a wrong result at worst. + + The condition is deliberately "all tensor operands", not "any". Passing int64 + *indices* alongside float data is normal and works -- ``embedding(idx, w)`` + and ``index_select`` are verified correct on device, and diverting them would + cost real performance. What fails is arithmetic carried out in int64, which is + exactly the all-int64 case: ``torch.diff`` on a ``torch.arange`` (int64 by + default) reaches the vendor's ``sub``, correct for fp32/int32 and broken for + int64 -- so the op cannot simply be excluded either. + """ + import torch + + seen = False + for value in list(args) + list(kwargs.values()): + items = value if isinstance(value, (list, tuple)) else (value,) + for item in items: + if isinstance(item, torch.Tensor): + if item.dtype != torch.int64: + return False + seen = True + return seen + + +def device_guarded_config(flag_gems): + """Return FlagGems' op table with every kernel wrapped in a device guard. + + FlagGems allocates its intermediates with ``device=input.device`` but + launches Triton on the *current* device. Call an op on a tensor living on + flagos:1 while flagos:0 is current and the launch reads across devices -- + the tops runtime rejects it ("DeviceId[1] of memory VA ... is not match for + DeviceId[0] of stream"), surfacing as ``topsErrorInvalidValue`` or a fault + rather than anything actionable. ``embedding`` is the case that shows up in + practice, through the ``index_select`` it decomposes to. + + The C++ dispatch path solves this with a ``c10::OptionalDeviceGuard`` per + caller (python_op_caller.cc), but ops that ``flag_gems.enable()`` registers + straight onto PrivateUse1 never pass through it. This is the same guard, in + Python, applied at the one place every such op goes through. + + The device is resolved from the first flagos tensor operand that names an + index -- matching the C++ DeviceOfArgs -- so leading non-tensor arguments + (``topk(values, k)``) resolve correctly. Ops with no such operand run + unguarded on the current device, as before. + + Pass the result as ``flag_gems.enable``'s config. It reads + ``flag_gems._FULL_CONFIG`` directly, so callers replace that attribute for + the duration of the call. + """ + import functools + + import torch + + from torch_fl import flagos + + def _first_device(args, kwargs): + for value in list(args) + list(kwargs.values()): + if isinstance(value, torch.Tensor): + if value.device.type in ("flagos", "privateuseone"): + return value.device.index + elif isinstance(value, (list, tuple)): + for item in value: + if isinstance(item, torch.Tensor) and item.device.type in ( + "flagos", + "privateuseone", + ): + return item.device.index + return None + + def _cpu_compute(aten_op, args, kwargs): + """Run an op on CPU operands and put flagos results back on the device. + + Used only for the int64 case below. An op registered into + ``TORCH_LIBRARY_IMPL(aten, PrivateUse1)`` cannot defer to the C++ kernel + for the same key -- returning NotImplemented re-enters this wrapper -- so + the fallback is performed here rather than delegated. + + ``aten_op`` is called rather than the FlagGems function: a vendor + override launches its Triton kernel regardless of where its operands + live, so handing it CPU tensors still compiles an int64 kernel and still + fails. Going through the aten op on CPU operands reaches the CPU kernel. + """ + import torch + + device = None + + def to_cpu(value): + nonlocal device + if isinstance(value, torch.Tensor): + if device is None and value.device.type in ("flagos", "privateuseone"): + device = value.device + return value.cpu() + if isinstance(value, (list, tuple)): + return type(value)(to_cpu(item) for item in value) + return value + + cpu_args = tuple(to_cpu(a) for a in args) + cpu_kwargs = {k: to_cpu(v) for k, v in kwargs.items()} + out = aten_op(*cpu_args, **cpu_kwargs) + if device is None: + return out + + def back(value): + if isinstance(value, torch.Tensor): + return value.to(device) + if isinstance(value, (list, tuple)): + return type(value)(back(item) for item in value) + return value + + return back(out) + + def _aten_op(aten_name): + """Resolve "sub.Tensor" to torch.ops.aten.sub.Tensor, or None.""" + import torch + + parts = str(aten_name).split(".") + op = getattr(torch.ops.aten, parts[0], None) + if op is None: + return None + if len(parts) > 1: + op = getattr(op, parts[1], None) + elif hasattr(op, "default"): + op = op.default + return op + + def _guard(func, aten_name): + # In-place and out= ops write through their operand, which a CPU copy + # would not propagate, so they keep the device path even for int64. + name = getattr(func, "__name__", "") + aten_op = None + if not (name.endswith("_") or name.endswith("_out")): + aten_op = _aten_op(aten_name) + + @functools.wraps(func) + def wrapper(*args, **kwargs): + if aten_op is not None and _is_int64_arithmetic(args, kwargs): + return _cpu_compute(aten_op, args, kwargs) + index = _first_device(args, kwargs) + if index is None or index == flagos.current_device(): + return func(*args, **kwargs) + with flagos.device(index): + return func(*args, **kwargs) + + return wrapper + + # functools.wraps carries __name__ and __module__ over, which the exclusion + # filter and the vendor-override detection both key on. + return tuple( + (entry[0], _guard(entry[1], entry[0])) + tuple(entry[2:]) + for entry in flag_gems._FULL_CONFIG + ) + + +def bind_vendor_ops_in_generic_modules(flag_gems) -> int: + """Make generic FlagGems ops call the vendor kernel for their sub-ops. + + A few generic ops are written on top of other ops, imported by value at + module scope -- ``flag_gems/ops/linear_backward.py`` does + ``from .mm import mm``, ``baddbmm.py`` does ``from .bmm import bmm``. That + binding is to the *generic* implementation and is fixed at import time, so + it survives the vendor's override: aten::mm dispatches to Enflame's + ``gcu300.ops.mm``, but ``linear_backward`` still calls ``flag_gems.ops.mm``. + + On GCU the two are not interchangeable. The generic ``mm_kernel_general`` + casts its index arithmetic with ``.to(tl.int64)`` (mm.py:101-104), and the + GCU Triton backend rejects the resulting widening -- "failed to legalize + operation 'arith.extsi' that was explicitly marked illegal", the same int64 + limitation as the rest of the tops stack. The vendor ``mm_kernel`` exists + precisely because of that, and stays on int32 indices. So every backward + pass through ``nn.Linear`` failed to compile while a direct ``torch.mm`` + worked, on the same shapes. + + Excluding ``linear_backward`` instead is not an option: aten has no CPU + kernel for it, so leaving it unregistered turns the failure into a hard + NotImplementedError rather than a fallback. + + The vendor overrides are discovered from ``flag_gems._FULL_CONFIG`` (the + op table FlagGems is about to register) by their module not being under + ``flag_gems`` -- the arch-specific package is imported under a bare name + like ``gcu300.ops.mm``. Only names a generic module imported from a sibling + generic op module are rebound; its own definitions are left alone. + + Call after ``import flag_gems``. Returns the number of names rebound. + """ + import sys + import types + + config = getattr(flag_gems, "_FULL_CONFIG", None) + if not config: + return 0 + + vendor = {} + for entry in config: + func = entry[1] if isinstance(entry, (tuple, list)) and len(entry) > 1 else None + if not isinstance(func, types.FunctionType): + continue + module = getattr(func, "__module__", "") or "" + if module == "flag_gems" or module.startswith("flag_gems."): + continue + vendor.setdefault(func.__name__, func) + if not vendor: + return 0 + + rebound = 0 + for name, module in list(sys.modules.items()): + if module is None or not name.startswith("flag_gems.ops."): + continue + for attr, value in list(vars(module).items()): + if not isinstance(value, types.FunctionType): + continue + # Only names imported from a *sibling* generic op module: an op's + # own definitions are what the vendor overrides, not what it calls. + origin = getattr(value, "__module__", "") or "" + if origin == name or not origin.startswith("flag_gems.ops."): + continue + replacement = vendor.get(value.__name__) + if replacement is None or replacement is value: + continue + setattr(module, attr, replacement) + rebound += 1 + return rebound diff --git a/torch_fl/flagos/__init__.py b/torch_fl/flagos/__init__.py index a47906cf..aac61a89 100644 --- a/torch_fl/flagos/__init__.py +++ b/torch_fl/flagos/__init__.py @@ -97,10 +97,17 @@ def _lazy_init(): # Eagerly import FlagGems to avoid deep import chain during dispatch. # FlagGems has a deep lazy import chain (fused → FLA → utils → models → sqlalchemy) # that can exceed Python's recursion limit when triggered inside PyTorch dispatch. + # + # Not just ImportError: FlagGems' vendor autodetect raises RuntimeError("No + # device were detected on your machine") when it is installed but cannot + # identify the backend -- which is the normal state on a vendor box whose + # Triton plugin is absent, and must not take down device init. The FlagGems + # kernels are optional everywhere; whatever is not registered runs on the + # native or CPU path. try: import flag_gems # noqa: F401 - except ImportError: - pass # FlagGems not installed, skip + except Exception: + pass # FlagGems unavailable or undetectable here, skip # Monkey-patch Tensor.__getitem__ to work around PyTorch C++ dispatch issue # with advanced indexing on custom devices. The C++ __getitem__ fails for @@ -266,6 +273,38 @@ def __new__(cls, device=None, priority=0, **kwargs): return super().__new__(cls, device=device, priority=priority, **kwargs) +class _DefaultStreamHandle: + """Minimal stream stand-in for backends with no CUDA runtime. + + Vendor Triton launchers (and FlagGems through them) want an object exposing + a raw stream handle, under whichever name their vendor uses -- ``cuda_stream`` + for CUDA-derived backends, ``gcu_stream`` for Enflame's triton_gcu. On + Enflame GCU, Ascend and MUSA the kernels go to the vendor's default stream, + which every one of them denotes with handle 0. The synchronization those + backends need already happens in their own op paths, so + ``synchronize``/``wait_stream`` have nothing to do here. + """ + + __slots__ = ("cuda_stream", "gcu_stream", "device_index") + + def __init__(self, device_index: int = 0): + self.cuda_stream = 0 + self.gcu_stream = 0 + self.device_index = device_index + + def __int__(self) -> int: + return 0 + + def synchronize(self): + pass + + def wait_stream(self, other): + pass + + def query(self) -> bool: + return True + + def _real_current_stream(device=None): """Resolve the actual current CUDA stream as a real ``torch.cuda.Stream``. @@ -276,8 +315,17 @@ def _real_current_stream(device=None): C++ runtime (``_cuda_getCurrentStream``) and rebuild a real Stream from it, so events record/wait on the same physical (default) stream the boxing kernels submit to. + + On a vendor backend with no CUDA runtime at all (Enflame GCU, Ascend, MUSA) + ``_cuda_getCurrentStream`` does not exist and ``torch.cuda`` cannot even be + lazily initialized ("Torch not compiled with CUDA enabled"). There is no CUDA + stream to describe, so return a stand-in carrying the vendor's default stream + handle: callers such as FlagGems' Triton launcher only read ``.cuda_stream`` + off the result, and those backends submit to their default stream. """ idx = current_device() if device is None else int(device) + if not hasattr(torch._C, "_cuda_getCurrentStream"): + return _DefaultStreamHandle(idx) stream_id, device_index, device_type = torch._C._cuda_getCurrentStream(idx) return torch.cuda.Stream( stream_id=stream_id, device_index=device_index, device_type=device_type