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

lxd-cumt · 2026-09-14T08:01:18Z

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.*).

…rectory coverage) This commit completes the CUDA plugin compatibility patches across the entire codebase: 1. Debug directory (transformer_engine/debug/): - fake_quant.py: Replace is_cuda checks with te_device_type() - log_fp8_tensor_stats.py: Replace torch.device('cuda') and torch.cuda.get_device_capability() - per_tensor_scaling.py: Replace is_cuda checks with te_device_type() - utils/stats_buffer.py: Replace .cuda() calls with .to(te_device_type()) 2. Additional missing files from cuda_patch_rc2.14: - permutation.py: Replace all is_cuda checks with device.type comparisons - tensor/float8_tensor.py: Replace is_cuda checks with device.type comparisons - tensor/float8_blockwise_tensor.py: Replace is_cuda checks with device.type comparisons - tensor/mxfp8_tensor.py: Replace is_cuda checks with device.type comparisons 3. Core infrastructure: - transformer_engine/__init__.py: Export te_platform and te_device_type from pytorch module - transformer_engine/pytorch/__init__.py: Define te_platform() and te_device_type() functions All runtime CUDA hardcodes have been replaced with plugin-compatible abstractions. Total: 64 files modified (+405/-289 lines) Signed-off-by: Xianduo Li <lixianduo@mail.nankai.edu.cn>

greptile-apps · 2026-09-14T08:04:38Z

Confidence Score: 4/5

[Medium impact] The changes reach many parts of Transformer Engine’s PyTorch workflow.

The single-file plugin fix looks safe, but vendor autocast on older supported PyTorch versions still needs fixing before merging.

Findings

Vendor autocast is ignored ▶

Summary

The PR adds configurable device helpers and replaces hardcoded CUDA device choices across the PyTorch code.

PyTorch plugins can set the device type Transformer Engine uses.
Modules, tensors, and device checks use the configured device type.
Autocast settings follow the configured device type.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A[Load native extension] --> B[Register original extension alias]
    B --> C[Import plugin]
    C --> D{Package with patches?}
    D -- Yes --> E[Apply patches]
    D -- No --> F[Load plugin]
    E --> F
    F --> G[Set device defaults and import PyTorch modules]

_{Reviews (18) · Last reviewed commit: "Skip patch discovery for single-file plu..." · Reviewed by Greptile}

The te_platform and te_device_type functions are dynamically assigned to the transformer_engine module namespace via: transformer_engine.te_platform = te_platform transformer_engine.te_platform = te_platform in transformer_engine/pytorch/__init__.py, so explicit imports in transformer_engine/__init__.py cause black to fail parsing due to circular import detection. Signed-off-by: Xianduo Li <lixianduo@mail.nankai.edu.cn>

Move te_platform and te_device_type imports outside of parenthesized import lists to fix SyntaxError. The imports were incorrectly placed inside multi-line import statements, preventing the modules from loading. Fixed files: - transformer_engine/debug/features/log_fp8_tensor_stats.py - transformer_engine/pytorch/attention/dot_product_attention/backends.py - transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py - transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py - transformer_engine/pytorch/attention/inference.py - transformer_engine/pytorch/graph.py - transformer_engine/pytorch/quantization.py - transformer_engine/pytorch/quantized_tensor.py - transformer_engine/pytorch/transformer.py Signed-off-by: Xianduo Li <lixianduo@mail.nankai.edu.cn>

Several files imported only te_platform but also called te_device_type(), causing NameError at runtime. Add te_device_type to imports. Fixed files: - transformer_engine/pytorch/attention/dot_product_attention/flex_attention.py - transformer_engine/pytorch/attention/rope.py - transformer_engine/pytorch/cpp_extensions/gemm.py - transformer_engine/pytorch/optimizers/fused_adam.py - transformer_engine/pytorch/utils.py Signed-off-by: Xianduo Li <lixianduo@mail.nankai.edu.cn>

Move te_platform and te_device_type imports outside of parenthesized import lists to fix SyntaxError across 20 more files. Fixed files: - transformer_engine/pytorch/attention/dot_product_attention/utils.py - transformer_engine/pytorch/attention/multi_head_attention.py - transformer_engine/pytorch/cpu_offload.py - transformer_engine/pytorch/module/base.py - transformer_engine/pytorch/module/grouped_linear.py - transformer_engine/pytorch/module/layernorm_linear.py - transformer_engine/pytorch/module/layernorm_mlp.py - transformer_engine/pytorch/module/linear.py - transformer_engine/pytorch/ops/basic/basic_linear.py - transformer_engine/pytorch/ops/basic/grouped_linear.py - transformer_engine/pytorch/ops/basic/l2normalization.py - transformer_engine/pytorch/ops/fused/forward_linear_scale_add.py - transformer_engine/pytorch/ops/fused/grouped_mlp.py - transformer_engine/pytorch/ops/fused/userbuffers_backward_linear.py - transformer_engine/pytorch/ops/fused/userbuffers_forward_linear.py - transformer_engine/pytorch/ops/op.py - transformer_engine/pytorch/tensor/nvfp4_tensor.py - transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py - transformer_engine/pytorch/triton/mhc.py - transformer_engine/pytorch/triton/permutation.py Signed-off-by: Xianduo Li <lixianduo@mail.nankai.edu.cn>

In make_row_id_map, use routing_map.device for tensor allocation instead of te_device_type() to respect the input tensor's device. Signed-off-by: Xianduo Li <lixianduo@mail.nankai.edu.cn>

Revert all te_platform() method calls back to torch.cuda.* calls. Keep te_device_type() string replacements for device type parameters. Preserve te_platform definition in transformer_engine/pytorch/__init__.py. Changes: - Reverted 146 te_platform() calls back to torch.cuda.* - Removed 'from transformer_engine import te_platform' imports - Kept all te_device_type() replacements intact - Kept te_platform definition in __init__.py for future use Signed-off-by: Xianduo Li <lixianduo@mail.nankai.edu.cn>

…istributed Replace remaining te_platform() calls with torch.cuda equivalents: - torch.cuda.nvtx.range() / range_push() / range_pop() for NVTX profiling - torch.cuda.amp.autocast for autocast context manager - Update docstring reference in fused_adam.py All te_platform() usage has been removed except for the definition in transformer_engine/pytorch/__init__.py. Signed-off-by: Xianduo Li <lixianduo@mail.nankai.edu.cn>

Signed-off-by: Xianduo Li <lixianduo@mail.nankai.edu.cn> # Conflicts: # transformer_engine/pytorch/module/layernorm_mlp.py

greptile-apps · 2026-09-22T03:01:47Z

Comments Outside Diff

These findings could not be posted inline.

Profiling bypasses configured platform transformer_engine/pytorch/utils.py:902 ▶

If NVTE_NVTX_ENABLED is enabled with a non-CUDA plugin, this code sends profiling ranges through torch.cuda.nvtx instead of the configured platform. The context, push, and pop helpers all do this, so module forward and backward paths can fail without CUDA or instrument the wrong backend.
Legacy autocast bypasses platform transformer_engine/pytorch/distributed.py:320 ▶

If a non-CUDA plugin is used with supported PyTorch 2.1–2.3, this branch enters torch.cuda.amp.autocast instead of the configured platform's autocast context. Checkpoint recomputation can therefore fail or use the wrong autocast behavior. The same regression appears in transformer_engine/pytorch/utils.py, where the legacy gpu_autocast_ctx is also bound to torch.cuda.amp.autocast.

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

ptrendx · 2026-10-09T23:14:15Z

@lxd-cumt Please fix the lint errors (not sure about those documentation errors, maybe they are also caused by the undefined variable?). Generally LGTM afterwards

Add the missing `te_device_type` import to eight modules that called it without importing it, which would raise NameError at runtime: jit, attention/inference, attention/dot_product_attention/{backends, context_parallel,utils}, ops/fused/{grouped_mlp, userbuffers_backward_linear}, and triton/mhc. Move the `te_device_type` import ahead of relative imports where pylint classifies it as a first-party import that must precede local ones. In jit and distributed, place it so it does not split the torch import group. Order `transformer_engine_torch` before `transformer_engine` in ops/_common. Add docstrings to `te_device_type` and `te_platform`. Signed-off-by: Xianduo Li <lixianduo@mail.nankai.edu.cn>

Resolve conflicts in two PyTorch modules where upstream touched the same lines as the device-type patch: - module/base.py: upstream reworked `get_dummy_wgrad` to add a ZeroTensor path and hoisted the target device into a local. Apply the device-type substitution to that local so the cache key, the ZeroTensor path, and the `torch.empty` call all share one device. - module/grouped_linear.py: upstream refactored the forward and backward implementations. Take the refactored file and re-apply the device-type substitution to the `m_splits` device check and the `device` constructor default and its docstring. Signed-off-by: Xianduo Li <lixianduo@mail.nankai.edu.cn>

greptile-apps · 2026-10-19T05:28:59Z

        dtype: torch.dtype
        if torch.is_autocast_enabled():
-            dtype = torch.get_autocast_dtype("cuda")
+            dtype = torch.get_autocast_dtype(te_device_type())


Vendor autocast is ignored

With a non-CUDA TE_DEVICE_TYPE, this code reads the configured device’s autocast dtype but still checks CUDA’s enabled state through torch.is_autocast_enabled(). When vendor autocast is on and CUDA autocast is off, the operation keeps the input dtype instead of the requested dtype.

Pass the configured device to both checks, with the needed version guard. The same mismatch appears in SwiGLU, grouped linear operations, and fused linear operations.

The device-type patch had changed the autocast dtype lookups to the configured device while leaving the enabled checks on CUDA. With a non-CUDA device type, vendor autocast enabled, and CUDA autocast disabled, these sites took the non-autocast branch and kept the input dtype instead of the requested autocast dtype. Route the operation sites through `maybe_autocast_dtype`, which already queries one device for both calls and carries the torch version guard. This covers the activation, SwiGLU, grouped linear, basic linear, and fused linear forward operations. Add `torch_is_autocast_gpu_enabled` next to `torch_get_autocast_gpu_dtype` for the two sites outside the operation layer that had the same mismatch: `set_activation_dtype`, which sets the activation dtype for every module, and the TransformerLayer forward cast. Signed-off-by: Xianduo Li <lixianduo@mail.nankai.edu.cn>

`tex.thd_get_partitioned_indices` allocates a CUDA output tensor and launches a CUDA kernel on the current CUDA stream, so it runs only on CUDA. The device-type patch had changed its dispatch condition to the configured device type, which sent packed sequences into that kernel on a non-CUDA device instead of the device-capable fallback, failing the `is_cuda` assertion before the fallback could run. Restore the CUDA condition in `build_rank_indices` and in `_get_thd_partition_cu_seqlens`, which had the same problem when normalizing the boundaries that feed the kernel, and record why both are gated on CUDA rather than on the configured device. Keep the configured device in the NO_LOAD_BALANCE branch, which builds indices with plain tensor operations and runs no CUDA kernel. Restore the assertion message, which named the device type that triggers it. Signed-off-by: Xianduo Li <lixianduo@mail.nankai.edu.cn>

Narrow the diff against main to the device-type substitutions themselves: - context_parallel.py: remove the comments on the THD partitioning device gates. - cpp_extensions/gemm.py: restore the single-line Float8BlockwiseQTensorStorage import, which did not need wrapping. - graph.py: remove an added blank line, leaving the file unchanged from main. Signed-off-by: Xianduo Li <lixianduo@mail.nankai.edu.cn>

Two paths reached from main used CUDA device checks where the work below them is not CUDA-specific: - tensor/nvfp4_tensor.py: `NVFP4Tensor` kept CUDA checks in its `record_stream` dispatch and in `copy_`, while the sibling `Float8Tensor`, `MXFP8Tensor`, and `Float8BlockwiseQTensor` classes use the configured device type for the same logic. On another device type `copy_` would move the operand to the destination device and `record_stream` would skip stream tracking. - cross_entropy.py: the input validation rejected non-CUDA tensors even though the Triton backend derives every allocation from the input device and contains nothing CUDA-specific. Leave the checks that guard compiled CUDA kernels, the `device_types` dispatch keys on custom operations, and the `is_cuda` properties that implement the PyTorch tensor contract. Signed-off-by: Xianduo Li <lixianduo@mail.nankai.edu.cn>

lxd-cumt · 2026-10-19T08:37:44Z

@lxd-cumt Please fix the lint errors (not sure about those documentation errors, maybe they are also caused by the undefined variable?). Generally LGTM afterwards

I have fixed the previous CI errors. Please approve and rerun CI. Thanks.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Xianduo Li <lixianduo@mail.nankai.edu.cn>

The top-level hook in transformer_engine/__init__.py imported <plugin>.patches before load_framework_extension() registered transformer_engine_torch_nv. A plugin that imports that alias at package level - a case the loader explicitly supports - failed there, and the enclosing except only warned, so apply_patches() never ran. The loader later imported the plugin successfully but never retried the patches, leaving device patches unapplied while vendor calls still went through CUDA. Move the call into load_framework_extension(), after the alias is registered and before the PyTorch submodules capture references to torch.cuda entry points. Probe the submodule with find_spec so a plugin that needs no patches may omit it, without masking real import failures inside an existing patches module. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Xianduo Li <lixianduo@mail.nankai.edu.cn>

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Xianduo Li <lixianduo@mail.nankai.edu.cn>

The modules whose `device` argument now defaults to `te_device_type()` allocate their parameters on the configured device, but the `prepare_forward()` input check still rejected anything whose `is_cuda` was false. That property is a tensor attribute, so patching the `torch.cuda` namespace does not change it, and forward raised before a plugin's native operation could run. Compare `device.type` against `te_device_type()` instead, and report the configured device type in the error message. One check covers Linear, LayerNormLinear, GroupedLinear, and LayerNormMLP, which share this base method. Behavior on CUDA is unchanged. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Xianduo Li <lixianduo@mail.nankai.edu.cn>

The remote branch carried an equivalent merge of upstream main. Both sides already contain 6b5cd7e, so this reconciles the histories without changing any file content. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Xianduo Li <lixianduo@mail.nankai.edu.cn>

Removing the plugin hook left an extra blank line behind. Drop it so the file matches main and no longer appears in the diff. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Xianduo Li <lixianduo@mail.nankai.edu.cn>

find_spec() raises ModuleNotFoundError instead of returning None when the parent module is not a package, so a single-file plugin that provides load_plugins() failed the optional-patches probe. The exception handler then disabled the plugin and restored the original extension. Check for __path__ before searching for the submodule. Genuine import failures inside an existing patches module still propagate. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Xianduo Li <lixianduo@mail.nankai.edu.cn>

lxd-cumt · 2026-10-21T08:13:16Z

Fixed the issues raised in the Greptile Robot review above.

lxd-cumt requested review from cyanguwa, ksivaman and timmoon10 as code owners September 14, 2026 08:01

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

lxd-cumt changed the title ~~fix: patch cuda hardcodes for plugin compatibility (complete debug di…~~ Abstract CUDA hardcodes into configurable te_device_type / te_platform Sep 14, 2026

greptile-apps Bot reviewed Sep 14, 2026

View reviewed changes

Comment thread transformer_engine/pytorch/attention/dot_product_attention/backends.py Outdated

Comment thread transformer_engine/pytorch/attention/dot_product_attention/flex_attention.py Outdated

lxd-cumt added 7 commits September 14, 2026 16:07

fix: use routing_map device instead of hardcoded device type

d7b7fe9

In make_row_id_map, use routing_map.device for tensor allocation instead of te_device_type() to respect the input tensor's device. Signed-off-by: Xianduo Li <lixianduo@mail.nankai.edu.cn>

lxd-cumt mentioned this pull request Sep 14, 2026

Abstract CUDA hardcodes into configurable te_device_type / te_platform #3113

Closed

Merge remote-tracking branch 'origin/main' into cuda_patch_main_latest

77f644d

Signed-off-by: Xianduo Li <lixianduo@mail.nankai.edu.cn> # Conflicts: # transformer_engine/pytorch/module/layernorm_mlp.py

lxd-cumt added 2 commits September 22, 2026 11:29

keep sync with torch.cuda api usage

1bad7ee

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

merge latest main and fix conflicts

a323d3c

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

ptrendx self-assigned this Sep 23, 2026

lxd-cumt added 3 commits October 8, 2026 14:52

merge latest main and fix conflicts

b946091

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

format

2d63707

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

merge latest main and fix conflicts

8fbfcfe

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

lxd-cumt added 2 commits October 19, 2026 13:09

greptile-apps Bot reviewed Oct 19, 2026

View reviewed changes

lxd-cumt added 2 commits October 19, 2026 14:19

lxd-cumt added 2 commits October 19, 2026 14:36

Merge latest main into cuda_patch_main_latest

3f1e95c

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Xianduo Li <lixianduo@mail.nankai.edu.cn>

greptile-apps Bot reviewed Oct 21, 2026

View reviewed changes

Comment thread transformer_engine/__init__.py Outdated

Merge branch 'main' into cuda_patch_main_latest

5f8b277

greptile-apps Bot reviewed Oct 21, 2026

View reviewed changes

Comment thread transformer_engine/pytorch/module/linear.py

lxd-cumt and others added 5 commits October 21, 2026 15:28

Merge latest main into cuda_patch_main_latest

a3f9403

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Xianduo Li <lixianduo@mail.nankai.edu.cn>

greptile-apps Bot reviewed Oct 21, 2026

View reviewed changes

Comment thread transformer_engine/common/__init__.py Outdated

Navigation Menu

Repository navigation

Abstract CUDA hardcodes into configurable te_device_type / te_platform - #3513

Abstract CUDA hardcodes into configurable te_device_type / te_platform#3513
lxd-cumt wants to merge 28 commits into
NVIDIA:mainfrom
lxd-cumt:cuda_patch_main_latest

lxd-cumt commented Sep 14, 2026 •

edited

Loading

Uh oh!

greptile-apps Bot commented Sep 14, 2026 •

edited

Loading

Uh oh!

Uh oh!

Uh oh!

greptile-apps Bot commented Sep 22, 2026 •

edited

Loading

Uh oh!

ptrendx commented Oct 9, 2026

Uh oh!

greptile-apps Bot Oct 19, 2026

Uh oh!

Uh oh!

lxd-cumt commented Oct 19, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

lxd-cumt commented Oct 21, 2026

Uh oh!

This branch has not been deployed

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

lxd-cumt commented Sep 14, 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 Sep 14, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Confidence Score: 4/5

Findings

Summary

Diagram

Uh oh!

Uh oh!

Uh oh!

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

Uh oh!

Comments Outside Diff

Uh oh!

ptrendx commented Oct 9, 2026

Uh oh!

greptile-apps Bot Oct 19, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

lxd-cumt commented Oct 19, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

lxd-cumt commented Oct 21, 2026

Uh oh!

This branch has not been deployed

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

lxd-cumt commented Sep 14, 2026 •

edited

Loading

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

greptile-apps Bot commented Sep 14, 2026 •

edited

Loading

greptile-apps Bot commented Sep 22, 2026 •

edited

Loading