Abstract CUDA hardcodes into configurable te_device_type / te_platform by lxd-cumt · Pull Request #3113 · NVIDIA/TransformerEngine

lxd-cumt · 2026-06-10T07:04:33Z

FlagOS Proposal: Plugin Architecture & Device-Agnostic Abstraction for TransformerEngine

Device-Type Abstraction: Replacing Hardcoded `"cuda"` References

The current TE PyTorch layer contains ~100 hardcoded "cuda" string literals and ~165 torch.cuda.* API calls. These span device placement (device="cuda"), autocast context (device_type="cuda"), device-type guards (device.type == "cuda"), and RNG state management (torch.cuda.CUDAGraph, torch.cuda._lazy_call). This makes TE non-functional on alternative accelerator platforms without invasive patching.

Proposed Design

Soft abstraction – A global te_device_type() / te_platform() accessor replaces ~200 literal "cuda" strings across the Python codebase.
Platform monkey-patch – A vendor-provided apply_patch() hook runs at import time to directly remap torch.cuda.* APIs (e.g. torch.cuda.device, torch.cuda.current_device, torch.cuda.current_stream) to the vendor equivalents (e.g. torch.other_vendor.*).

greptile-apps · 2026-06-10T07:12:42Z

Greptile Summary

This PR introduces a device-agnostic abstraction layer over the ~200 hardcoded "cuda" string literals and torch.cuda.* calls across the TransformerEngine PyTorch backend, replacing them with te_device_type() / te_platform() accessors and a plugin hook (NVTE_PLUGIN env var) that lets vendors monkey-patch CUDA APIs at import time.

Plugin system (transformer_engine/__init__.py): Loads a vendor-supplied {plugin}.patches.apply_patches() before the PyTorch backend is imported; failures surface as RuntimeWarning.
Accessor definitions (pytorch/__init__.py): TE_DEVICE_TYPE / TE_PLATFORM globals default to "cuda" / torch.cuda, overridable by a plugin; te_device_type() and te_platform() are injected back onto the top-level package so all 44 changed files can import them from transformer_engine.
Mechanical replacement (remaining 42 files): "cuda" string literals, tensor.is_cuda, device="cuda", and torch.get_autocast_dtype("cuda") are replaced with their te_device_type() / te_platform() equivalents across attention, tensor, optimizer, and op modules.

Confidence Score: 3/5

Several files have incomplete or incorrect migrations that will misplace tensors or misreport device state on non-CUDA accelerators, including a multi-GPU device-index bug in the Triton permutation kernel.
The grouped_tensor_storage.py graph-capture optimisation branch now checks te_device_type() for the device but immediately delegates to torch.cuda.is_available() / torch.cuda.is_current_stream_capturing(), so the branch behaves incorrectly on MUSA. nvfp4_tensor.py's get_rht_matrix appends .to(te_device_type()) after building all tensors on the caller-supplied device, silently overriding the device index and producing wrong placements on multi-GPU CUDA setups. triton/permutation.py's workspace_tensor uses bare te_device_type() while row_id_map follows routing_map.device, causing a device mismatch that triggers an illegal-memory-access on any non-default device. jit.py warmup functions had their tensor allocations updated but left torch.cuda.get/set_rng_state and torch.cuda.empty_cache untouched, breaking RNG reproducibility on MUSA hosts that rely solely on the explicit accessor rather than the monkey-patch.
transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py, transformer_engine/pytorch/tensor/nvfp4_tensor.py, transformer_engine/pytorch/triton/permutation.py, and transformer_engine/pytorch/jit.py each have concrete correctness issues that need to be resolved before this lands.

Important Files Changed

Filename	Overview
transformer_engine/init.py	Adds `NVTE_PLUGIN` env-var hook: dynamically imports a plugin module's `apply_patches()` before loading the PyTorch backend. Plugin failures now emit a `RuntimeWarning` instead of silently passing. Clean change.
transformer_engine/pytorch/init.py	Defines `te_device_type()` / `te_platform()` accessors and registers them on the top-level `transformer_engine` module. `gpu_autocast_ctx` and several function default-parameter values capture `te_device_type()` at module-load time; functionally correct given plugin ordering but fragile.
transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py	Incomplete migration: device-type check uses `te_device_type()` but the subsequent `torch.cuda.is_available()` and `torch.cuda.is_current_stream_capturing()` guards remain hardcoded, breaking the graph-capture optimisation path on non-CUDA accelerators.
transformer_engine/pytorch/tensor/nvfp4_tensor.py	`get_rht_matrix` appends `.to(te_device_type())` after `.to(dtype=torch.bfloat16)`, silently overriding the explicit `device: int` parameter and causing wrong device placement when the caller's device differs from the default device.
transformer_engine/pytorch/triton/permutation.py	`make_row_id_map` correctly updates `row_id_map` to use `routing_map.device`, but `workspace_tensor` still uses bare `te_device_type()` (no device index), causing a device mismatch when the input is on a non-default device.
transformer_engine/pytorch/jit.py	Tensor allocations updated to `te_device_type()`, but `torch.cuda.get_rng_state/set_rng_state/empty_cache` remain hardcoded in all three warmup functions, breaking RNG reproducibility on non-CUDA accelerators.
transformer_engine/pytorch/ops/fused/userbuffers_backward_linear.py	Device check updated to `te_device_type()`; error message has a double space before "devices" (cosmetic). Import guard for `bulk_overlap_ag_with_external_gemm` was a pre-existing issue noted in prior review.
transformer_engine/pytorch/utils.py	`canonicalize_device` and `devices_match` correctly use `te_platform().current_device()`. `gpu_autocast_ctx` captures `te_device_type()` at module load time via `functools.partial`, which is fine given plugin-before-import ordering.
transformer_engine/pytorch/permutation.py	All `is_cuda` checks replaced with `device.type == te_device_type()` and error messages updated consistently. No issues found.
transformer_engine/pytorch/module/base.py	Platform availability check and device guards updated to `te_platform()` / `te_device_type()`. Uses `raise RuntimeError` for the availability check (correct pattern). `torch.load(..., map_location=te_device_type())` is valid.

Sequence Diagram

sequenceDiagram
    participant User
    participant TE as transformer_engine.__init__
    participant Plugin as vendor_plugin.patches
    participant PYINIT as pytorch.__init__
    participant Module as TE Module (e.g. backends.py)

    User->>TE: import transformer_engine
    TE->>TE: read NVTE_PLUGIN env var
    alt NVTE_PLUGIN set
        TE->>Plugin: import_module(plugin).apply_patches()
        Plugin->>TE: "set TE_DEVICE_TYPE = "musa""
        Plugin->>TE: "set TE_PLATFORM = torch.musa"
        Plugin-->>TE: "patch torch.cuda.* → torch.musa.*"
    end
    TE->>PYINIT: from . import pytorch
    PYINIT->>TE: "TE_DEVICE_TYPE = getattr(te, "TE_DEVICE_TYPE", "cuda")"
    PYINIT->>PYINIT: define te_device_type() / te_platform()
    PYINIT->>TE: "te.te_device_type = te_device_type"
    PYINIT->>Module: import submodules
    Module->>TE: from transformer_engine import te_device_type
    Module->>Module: use te_device_type() / te_platform() at runtime

_{Reviews (7): Last reviewed commit: "Simplify plugin device/platform interfac..." | Re-trigger Greptile}

greptile-apps · 2026-06-10T07:12:45Z

+dpa_utils._original_get_attention_backend = dpa_utils.get_attention_backend
+# Replace dpa_utils.get_attention_backend with tex.get_attention_backend
+# This allows each backend (FlagOS, CUDA, Reference) to control its own backend selection
+dpa_utils.get_attention_backend = tex.get_attention_backend


AttributeError at import time breaks the PyTorch module

tex is transformer_engine_torch (the C++ extension), which does not expose a get_attention_backend attribute. Accessing tex.get_attention_backend directly on line 75 (without a getattr guard) raises AttributeError the moment transformer_engine.pytorch is imported, making the entire PyTorch backend unusable on any standard CUDA installation. The previous line (69) correctly uses getattr(tex, "flash_attention", _FlashAttentionNative) with a fallback — the same pattern must be applied here, or the unconditional attribute access must be removed.

greptile-apps · 2026-06-10T07:12:46Z

+try:
+    from transformer_engine_torch import bulk_overlap_ag_with_external_gemm
+except ImportError:
+    bulk_overlap_ag_with_external_gemm = None


NoneType call crash in backward pass when bulk_overlap_ag_with_external_gemm is unavailable

The import is now guarded (= None on failure), but line 435 calls bulk_overlap_ag_with_external_gemm(ub_obj_overlap_wgrad, dgrad_send_stream, dgrad_recv_stream) unconditionally. Whenever this code path is hit in a tensor-parallel row-overlap backward pass on a system where this symbol is absent, a TypeError: 'NoneType' object is not callable is raised at runtime rather than at import time. A guard like if bulk_overlap_ag_with_external_gemm is not None: (or raising a descriptive error earlier) is needed at the call site.

greptile-apps · 2026-06-10T07:12:47Z

+try:
+    from .plugin.core.backends.vendor.musa.patches import apply_patch as _musa_apply_patch
+
+    _musa_apply_patch()
+except Exception as e:
+    pass


Silent except Exception: pass hides all MUSA patch failures

The bare except Exception as e: pass swallows every failure during the MUSA patch import — including AttributeError raised when torch.musa.* attributes referenced in _PATCH_CALLS don't exist on standard CUDA systems. The variable e is never logged or inspected. On CUDA systems this is the common path (no torch.musa), so every import of transformer_engine silently triggers and discards an exception. At minimum, emit a logging.debug or use a narrower exception type (e.g., ImportError) and let other errors propagate.

greptile-apps · 2026-06-10T07:12:48Z

+# Patches: (parent_object, attribute_name, replacement_callable)
+_PATCH_CALLS: list[tuple[object, str, Callable[..., object]]] = [
+    # We do not recommend replace is_available, due to its device-related behavior.
+    # (torch.cuda, "is_available", torch.musa.is_available),
+    (torch.cuda, "get_device_properties", torch.musa.get_device_properties),
+    (torch.cuda, "device", torch.musa.device),
+    (torch.cuda, "current_device", torch.musa.current_device),
+    (torch.cuda, "synchronize", torch.musa.synchronize),
+    (torch.cuda, "is_current_stream_capturing", torch.musa.is_current_stream_capturing),
+    # TODO: Add NVTX patches for MUSA.
+    # NVTX is CUDA-specific; make it a no-op on MUSA.
+    (torch.cuda.nvtx, "range_push", _noop),
+    (torch.cuda.nvtx, "range_pop", _noop),


_PATCH_CALLS accesses torch.musa.* at module load time

_PATCH_CALLS is a module-level list that dereferences torch.musa.get_device_properties, torch.musa.device, etc. when patches.py is imported. On any system without torch_musa, this raises AttributeError the moment the import is attempted. The caller in __init__.py wraps this in a blanket except Exception: pass, so it fails silently, but it means the module is broken by construction on non-MUSA hosts. Deferring the attribute lookups to inside apply_patch() (where hasattr(torch, "musa") is already checked) would make the module safe to import on all platforms.

greptile-apps · 2026-06-10T07:12:51Z

+    # Mark TE global device type for Python-side callers.
+    # IMPORTANT: do not import `transformer_engine` here, because TE's `__init__.py`
+    # imports this module to run patches and that would cause a circular import.
+    try:
+        import transformer_engine
+
+        transformer_engine.TE_DEVICE_TYPE = "musa"
+        transformer_engine.TE_PLATFORM = torch.musa


The comment warns against importing transformer_engine here due to a circular-import risk, but the very next lines do exactly that. During transformer_engine/__init__.py execution Python's import cache returns the partially-initialized module, so TE_DEVICE_TYPE (set before the patch call) is reachable and the assignment works — but the comment creates a false sense of safety and the approach is still fragile if the import order ever changes.

Suggested change

# Mark TE global device type for Python-side callers.

# IMPORTANT: do not import `transformer_engine` here, because TE's `__init__.py`

# imports this module to run patches and that would cause a circular import.

try:

import transformer_engine

transformer_engine.TE_DEVICE_TYPE = "musa"

transformer_engine.TE_PLATFORM = torch.musa

# Mark TE global device type for Python-side callers.

# NOTE: importing `transformer_engine` here re-enters a partially-initialised module

# (its __init__.py is still running), but Python's import cache makes this safe as long

# as TE_DEVICE_TYPE is assigned before apply_patch() is called.

try:

import transformer_engine

transformer_engine.TE_DEVICE_TYPE = "musa"

transformer_engine.TE_PLATFORM = torch.musa

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

lxd-cumt · 2026-06-17T07:52:15Z

Thanks for the review! I've addressed the above Greptile comments.

Signed-off-by: Xianduo Li <lixianduo@mail.nankai.edu.cn>

…th explicit raise Signed-off-by: Xianduo Li <lixianduo@mail.nankai.edu.cn>

Signed-off-by: Xianduo Li <lixianduo@mail.nankai.edu.cn>

lxd-cumt · 2026-06-17T08:05:35Z

Should this PR be merged into the main branch, so it can follow future NVIDIA TE releases?

ptrendx

Left some comments. Could we use this PR also as an opportunity to take a look at each of the cases where we hardcoded the device type and see if we can instead use a device from e.g. the input tensor instead?

…edundant code - Move TE_DEVICE_TYPE, te_device_type, TE_PLATFORM, te_platform from top-level __init__.py to transformer_engine/pytorch/__init__.py since they depend on torch and the top-level package is shared with JAX. - Remove unnecessary rng_state patch in context_parallel.py - Remove unnecessary fused_attention thd check in utils.py - Remove extra blank line in linear.py - Remove redundant device-is-None check in float8_blockwise_tensor.py Signed-off-by: Xianduo Li <lixianduo@mail.nankai.edu.cn>

Replace NVTE_ENABLE_PLUGIN=1 with NVTE_PLUGIN=<module_name> so the plugin module is not hardcoded. Any plugin implementing <module>.patches.apply_patches() can be loaded via this env var. Signed-off-by: Xianduo Li <lixianduo@mail.nankai.edu.cn>

lxd-cumt · 2026-07-07T07:33:26Z

Left some comments. Could we use this PR also as an opportunity to take a look at each of the cases where we hardcoded the device type and see if we can instead use a device from e.g. the input tensor instead?

Thanks for the review, I have addressed the comments.
This PR touches similar hardcoded device types in TransformerEngine, so we can align on a technical approach together. And if we plan to upstream this to main eventually, we'll need to handle cuda hardcoding across more files as well.

ptrendx · 2026-07-29T09:09:41Z

+# Plugin system: set NVTE_PLUGIN to the plugin module name to enable.
+# e.g. NVTE_PLUGIN=transformer_engine_plugin_fl
+_nvte_plugin = os.environ.get("NVTE_PLUGIN")
+if _nvte_plugin:
+    try:
+        from importlib import import_module
+
+        _patches = import_module(f"{_nvte_plugin}.patches")
+        _patches.apply_patches()
+    except Exception as e:
+        import warnings
+
+        warnings.warn(
+            f"NVTE_PLUGIN={_nvte_plugin} but plugin patch apply failed: {e}",
+            RuntimeWarning,
+            stacklevel=1,
+        )


Is this needed in addition to #3107 or should it be removed once that PR is merged?

Previous pr #3401 replaces ops within transformer_engine_torch, whereas this patch targets torch.cuda-related APIs repo-wide. Despite the similar plugin load mechanism, they address orthogonal concerns. I suggest retaining this.

ptrendx · 2026-07-29T09:13:02Z

+    try:
+        return transformer_engine.TE_DEVICE_TYPE
+    except Exception:
+        return default


Considering that we set this value just above, why do we need try/except here? And if there was some
chance that this value got deleted somehow, then it should probably be an error anyway rather than
a quiet ignore.

ptrendx · 2026-07-29T09:13:17Z

+    try:
+        return transformer_engine.TE_PLATFORM
+    except Exception:
+        return default


Same comment here.

ptrendx · 2026-07-29T09:16:53Z

+if not hasattr(transformer_engine, "TE_DEVICE_TYPE"):
+    transformer_engine.TE_DEVICE_TYPE = "cuda"
+TE_DEVICE_TYPE = transformer_engine.TE_DEVICE_TYPE


Actually I'm not sure why this code lives here. We should be able to just set TE_DEVICE_TYPE (and
platform) inside transformer_engine/__init__.py before loading of the patches. Also, do we actually
need this very generic patches interface or can we just set those values directly (plugin would just

TE_PLATFORM defaults to torch.cuda, which requires torch to be importable. The top-level transformer_engine/__init__.py is shared with
JAX and does not import torch, so both TE_DEVICE_TYPE and TE_PLATFORM defaults are set here to keep them co-located. The functions
te_device_type() and te_platform() are then registered onto the top-level package for use by other submodules.

Regarding the generic patches interface: apply_patches() does more than setting these two values — it also monkey-patches torch.cuda.*
APIs (e.g. current_device, synchronize, get_device_properties) for non-CUDA backends. A direct-set approach would not cover those.

ptrendx · 2026-07-29T09:17:41Z

+        if is_training and device_compute_capability >= (10, 0):
+            logger.debug("Disabling FusedAttention for determinism reasons on Blackwell")
+            use_fused_attention = False
+            fused_attention_backend = None


This looks unrelated.

ptrendx · 2026-07-29T09:18:51Z

+        _alibi_cache["_alibi_bias"] = bias.contiguous().to(
+            dtype=bias_dtype, device=te_device_type()
+        )


Here the bias should already be on the device, so probably we could just omit the device parameter
here entirely.

ptrendx · 2026-07-29T09:25:25Z

    """
-    row_id_map = torch.empty((num_tokens, num_experts * 2 + 1), dtype=torch.int32, device="cuda")
+    row_id_map = torch.empty(
+        (num_tokens, num_experts * 2 + 1), dtype=torch.int32, device=te_device_type()


Use routing_map's device.

ptrendx · 2026-07-29T09:26:30Z

-        if device.type != "cuda":
-            device = torch.device("cuda", torch.cuda.current_device())
+        if device.type != te_device_type():
+            device = torch.device(te_device_type(),te_platform().current_device())


Missing space after comma.

- Remove unnecessary try/except in te_device_type() and te_platform() - Simplify TE_DEVICE_TYPE/TE_PLATFORM defaults using getattr - Remove unrelated Blackwell FusedAttention change - Use routing_map.device instead of te_device_type() in permutation - Fix missing space after comma in utils.py Signed-off-by: Xianduo Li <lixianduo@mail.nankai.edu.cn>

lxd-cumt · 2026-08-19T09:37:38Z

A few questions for discussion:

Target branch: Should this PR target main instead of release/v2.14?
te_device_type / te_platform abstraction: Is this approach acceptable? This implies that future development on main would need to use te_device_type() / te_platform() in place of hardcoded "cuda" hardstrings.
Megatron-LM: If this approach is agreed upon, I plan to open a similar PR for Megatron-LM to abstract device-specific references there as well.

github-actions Bot added the community-contribution PRs from external contributor outside the core maintainers, representing community-driven work. label Jun 10, 2026

greptile-apps Bot reviewed Jun 10, 2026

View reviewed changes

Comment thread transformer_engine/pytorch/module/base.py Outdated

lxd-cumt and others added 5 commits June 17, 2026 15:53

fake: try to patch cuda hardcodes

cdefc94

Signed-off-by: Xianduo Li <lixianduo@mail.nankai.edu.cn>

remove vendor related

e5e85a9

Signed-off-by: Xianduo Li <lixianduo@mail.nankai.edu.cn>

support apply paatches

aff2e8a

Signed-off-by: Xianduo Li <lixianduo@mail.nankai.edu.cn>

fix: use transformer_engine_plugin_fl interface and replace assert wi…

561efb3

…th explicit raise Signed-off-by: Xianduo Li <lixianduo@mail.nankai.edu.cn>

fix: patch cuda hardcodes for plugin compatibility

8ae24c3

Signed-off-by: Xianduo Li <lixianduo@mail.nankai.edu.cn>

lxd-cumt force-pushed the cuda_patch branch from 523df5a to 8ae24c3 Compare June 17, 2026 07:53

ptrendx reviewed Jun 30, 2026

View reviewed changes

Xianduo Li added 2 commits July 7, 2026 14:21

ptrendx reviewed Jul 29, 2026

View reviewed changes

ptrendx self-assigned this Aug 4, 2026

lxd-cumt force-pushed the cuda_patch branch from 4905215 to bccb702 Compare August 19, 2026 09:36

Uh oh!

Conversation

lxd-cumt commented Jun 10, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

FlagOS Proposal: Plugin Architecture & Device-Agnostic Abstraction for TransformerEngine

Device-Type Abstraction: Replacing Hardcoded "cuda" References

Proposed Design

Uh oh!

greptile-apps Bot commented Jun 10, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Greptile Summary

Confidence Score: 3/5

Important Files Changed

Sequence Diagram

Uh oh!

greptile-apps Bot Jun 10, 2026

Choose a reason for hiding this comment

Uh oh!

greptile-apps Bot Jun 10, 2026

Choose a reason for hiding this comment

Uh oh!

greptile-apps Bot Jun 10, 2026

Choose a reason for hiding this comment

Uh oh!

greptile-apps Bot Jun 10, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

greptile-apps Bot Jun 10, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

lxd-cumt commented Jun 17, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

lxd-cumt commented Jun 17, 2026

Uh oh!

ptrendx left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

lxd-cumt commented Jul 7, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

lxd-cumt commented Jun 10, 2026 •

edited

Loading

Device-Type Abstraction: Replacing Hardcoded `"cuda"` References

greptile-apps Bot commented Jun 10, 2026 •

edited

Loading

lxd-cumt commented Jun 17, 2026 •

edited

Loading