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