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 two-layer device-agnostic abstraction for TransformerEngine's PyTorch backend, replacing ~100 hardcoded \"cuda\" string literals with a te_device_type() accessor and ~165 torch.cuda.* API calls with a monkey-patchable te_platform() handle, enabling third-party accelerator plugins (e.g. MUSA) to remap both without modifying TE source.

A NVTE_PLUGIN environment variable loads a vendor plugin at transformer_engine/__init__.py import time; the plugin's apply_patches() sets transformer_engine.TE_DEVICE_TYPE and transformer_engine.TE_PLATFORM before transformer_engine.pytorch is ever imported.
transformer_engine/pytorch/__init__.py defines the te_device_type() / te_platform() accessor functions (defaulting to \"cuda\" / torch.cuda) and exports them onto the transformer_engine namespace for use across all 44 changed files.
The torch.cuda.* API calls that remain in the codebase (e.g. torch.cuda.get_rng_state, torch.cuda.default_generators) are intentionally left for the monkey-patch layer rather than the string-substitution layer.

Confidence Score: 3/5

Multiple issues flagged in earlier review rounds remain unaddressed on critical training paths, including a NoneType callable crash in the tensor-parallel backward path and an AttributeError at import time.

The string-substitution layer itself is mechanically correct across all 44 files and the plugin loading infrastructure is sound, but unresolved bugs in userbuffers_backward_linear.py and dot_product_attention.py affect core forward/backward execution paths.

transformer_engine/pytorch/ops/fused/userbuffers_backward_linear.py (unconditional call to possibly-None symbol), transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py (unguarded tex attribute access at import)

Important Files Changed

Filename	Overview
transformer_engine/init.py	Adds plugin loading via NVTE_PLUGIN env var; uses warnings.warn for failures, improving on the previous bare except: pass pattern flagged in prior review
transformer_engine/pytorch/init.py	Introduces TE_DEVICE_TYPE / TE_PLATFORM globals and te_device_type() / te_platform() accessors; TE_DEVICE_TYPE module-level alias is a frozen copy of the value at import time
transformer_engine/pytorch/utils.py	Replaces hardcoded 'cuda' strings with te_device_type()/te_platform() throughout; gpu_autocast_ctx partial captures device_type at import time rather than lazily
transformer_engine/pytorch/distributed.py	Default parameter of _get_cuda_rng_state updated to te_device_type(); function body still calls torch.cuda.current_device() and torch.cuda.default_generators directly (relies on monkey-patching per PR design)
transformer_engine/pytorch/module/base.py	Device availability check changed to te_platform().is_available(); input device guard updated; parameter init device string abstracted
transformer_engine/pytorch/permutation.py	All .is_cuda guards replaced with .device.type == te_device_type(); one error message uses string concatenation instead of f-string
transformer_engine/pytorch/ops/fused/userbuffers_backward_linear.py	Device guard updated to te_device_type(); error message has a double-space typo; bulk_overlap_ag_with_external_gemm still called unconditionally
transformer_engine/pytorch/ops/fused/userbuffers_forward_linear.py	Device guard and autocast dtype lookup updated to use te_device_type()
transformer_engine/pytorch/quantization.py	Default device in RecipeState subclasses changed from torch.device('cuda') to torch.device(te_device_type()); straightforward substitution
transformer_engine/pytorch/jit.py	Tensor allocations in JIT warmup functions updated; torch.cuda.get_rng_state() calls remain (relies on monkey-patching per PR design)
transformer_engine/pytorch/triton/permutation.py	All hardcoded 'cuda' device strings in tensor allocations replaced with te_device_type()
transformer_engine/pytorch/tensor/float8_tensor.py	Default device in quantizer and device detection logic updated to use te_device_type()

_{Reviews (5): Last reviewed commit: "fix: use NVTE_PLUGIN env var instead of ..." | 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!

greptile-apps · 2026-06-10T09:27:13Z

        super().__init__()
-        if not torch.cuda.is_available():
-            raise RuntimeError("TransformerEngine needs CUDA.")
+        assert te_platform().is_available(), f"TransformerEngine needs {te_device_type()}."


assert for runtime environment checks is stripped when Python runs with the -O flag, silently bypassing the guard. The original code used an explicit raise RuntimeError, which fires unconditionally. Replacing it with assert means a user running python -O can instantiate a TransformerEngineBaseModule on a system without the required hardware, only to get a confusing crash later deep inside a CUDA/MUSA kernel.

Suggested change

assert te_platform().is_available(), f"TransformerEngine needs {te_device_type()}."

if not te_platform().is_available():

raise RuntimeError(f"TransformerEngine needs {te_device_type()}.")

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?

ptrendx · 2026-06-30T21:10:01Z

 from importlib import metadata
 import transformer_engine.common

+import torch


We cannot import torch in this file - it is used both by pyTorch and JAX backends.

It should go into transformer_engine/pytorch/__init__.py instead.

fixed, thanks

ptrendx · 2026-06-30T21:15:20Z

+        # Fix: flash-attn 2.3.x ~ 2.6.x also needs rng_state for dropout
+        if not use_flash_attn_3 and rng_states is not None:
+            fa_backward_kwargs["rng_state"] = rng_states[cp_size - step - 1]


This change looks not connected to the rest of the PR - is it intended?

ptrendx · 2026-06-30T21:18:12Z

        if use_flash_attention:
            use_flash_attention = False
            logger.debug("Disabling FlashAttention for max_logit")
+        if use_fused_attention and qkv_format == "thd":


Changes in this file also seem disconnected from the rest of this PR and should be in its own PR.

ptrendx · 2026-06-30T21:20:35Z

 )
 from ...debug.pytorch.debug_state import TEDebugState

+


Unnecessary change.

ptrendx · 2026-06-30T21:22:56Z

+        if device is None:
+            device = torch.device(te_device_type())


This is already handled by the lines below.

…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.

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

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

	assert te_platform().is_available(), f"TransformerEngine needs {te_device_type()}."
	if not te_platform().is_available():
	raise RuntimeError(f"TransformerEngine needs {te_device_type()}.")

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

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!

greptile-apps Bot Jun 10, 2026

Choose a reason for hiding this comment

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!

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!

lxd-cumt commented Jul 7, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

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