[VenusTZZ]

Background

DeepSpeed has already embraced Triton for a growing set of fused kernels — the recently developed AutoEP (Automatic Expert Parallel) features, for example, rely on Triton kernels for token permutation and row weighting, alongside the grouped GEMM (General Matrix Multiply), SwiGLU, fused RMSNorm, and fused RoPE paths. As with Triton itself, that integration has so far been CUDA-first: the guards and tests naturally reference CUDA, the reference platform. In the meantime, the Triton ecosystem has expanded well beyond CUDA — several hardware platforms now ship their own Triton backends, for example the triton-ascend backend for Ascend NPUs — and a backend that reports itself through the standard driver probe can run the same kernels. This PR carries DeepSpeed's Triton integration toward accelerator-agnosticism: the device guards, launch contexts, and tests stop naming a specific device, and the Triton gate follows the shared driver probe, so every platform can benefit wherever its backend reports support. CUDA remains the primary, fully supported platform, and every default path behaves exactly as before.

Summary

Removes the CUDA hard-coding from the triton_ops device guards, launch contexts, and tests, and routes the Triton gate through the shared Triton driver probe so the fused kernels can run on any accelerator that reports a working Triton backend (CUDA first, others — e.g. NPU via triton-ascend — where their backend is installed). With no Triton backend installed, is_triton_available() is False and every gate stays closed — no behavior changes on any platform; this is a guard refactor, not a kernel change.

Problem

Hard-coded CUDA assumptions in deepspeed/ops/triton_ops

1. Device guards hard-code CUDA. fused_rms_norm.py, fused_rotary_pos_emb.py, and autoep_fused_token_ops.py reject any tensor with device.type != "cuda" ("runs CUDA kernels"), even though the kernels are written in Triton and the runtime device is whatever the accelerator reports.
2. Launch contexts hard-code torch.cuda.device. swiglu_triton.py (packed forward/backward) and fused_rotary_pos_emb.py switch devices with torch.cuda.device(x.device), which breaks on CPU-built torch and on non-CUDA accelerators.
3. The NPU gate is permanently closed. NPU_Accelerator.is_triton_supported() returns False unconditionally, so all Triton paths fall back on NPU even when a working Triton backend (triton-ascend) is installed.
4. The unit tests gate on CUDA by name. Module-level guards use device_name() == "cuda"; bfloat16 capability is queried via torch.cuda.get_device_capability() (crashes on CPU-built torch); guard-message assertions expect the "CUDA kernels" wording; two-device stream regressions use torch.cuda.device contexts; the sm-capability auto-selection test has no device guard at all.

Fix

• npu_accelerator.py: is_triton_supported() delegates to is_triton_available() — the same shared probe as the other accelerators, which reports True only when a Triton backend's driver is active. Without triton-ascend the gate stays closed, preserving today's fallbacks.
• _triton.py: adds is_accelerator_tensor(), delegating to the accelerator's on_accelerator(). This is the single place in the package that lazily imports get_accelerator, keeping the module import-light.
• fused_rms_norm.py / fused_rotary_pos_emb.py / autoep_fused_token_ops.py: device guards use the helper; error messages say "runs Triton kernels"; the "same CUDA device" message becomes "same device".
• swiglu_triton.py / fused_rotary_pos_emb.py: torch.cuda.device(...) launch contexts become device-agnostic torch.device(...) contexts.
• ep_kernels.py: permute_rows_supported accepts any accelerator tensor; docstrings say "accelerator tensor" instead of "CUDA tensor".
• Tests (tests/unit/v1/ops/triton_ops/): module-level gates use is_triton_supported(); bfloat16 checks use is_bf16_supported(); guard-message assertions expect the Triton wording; two-device contexts use torch.device(accelerator.device_name(...)); the sm-capability auto-selection test is marked CUDA-only.

Compliance with the repo AI-contribution policy (AGENTS.md)

• Duplicate-work check: searched open PRs/issues for is_triton_supported + NPU and triton + ascend — no existing PR or issue covers opening the NPU Triton gates or device-agnosticizing the triton_ops guards. The nearest PR, AutoEP: generate large NPU permutation indices on device #8822 (AutoEP NPU permutation indices), is closed and unmerged and touches a different path.
• Numerics / performance / memory impact: none — this is a gating/refactor change; every default path is unchanged.
• Compatibility matrix: not applicable — no training-loop numeric or memory behavior changes; validation is at the kernel-unit-test level only.
• AI assistance: this PR was prepared with the assistance of AI coding agents; the submitting human has reviewed every changed line and run the tests below.

Test plan

The affected suites run on both platforms against branch npu-triton-support (CUDA: A100; NPU: Ascend 910B4, triton-ascend 3.6.0).

CUDA (A100) — the primary platform

• triton_ops (305): all pass — identical to master.
• MoE AutoEP (permute_rows / unit / fused_parity / grad_parity, 259): 249 pass / 10 expected skips (CUDA-only DeepEP/H100 cases), no failures.
• test_chunked_cross_entropy.py (50): all pass — including the device-agnostic Triton-backend change.
• Total: 604 pass / 0 fail.

Environment: 2 × NVIDIA A100-SXM4-80GB (80 GB each), driver 580.126.16, CUDA 12.4, NCCL 2.21.5, PyTorch 2.14.1+cu126, Triton 3.8.0, transformers 5.20.2, Python 3.10.22.

NPU (Ascend 910B4, triton-ascend 3.6.0)

• triton_ops (305): 292 pass / 11 fail / 2 skip — the 11 failures are triton-ascend backend issues that have been reported upstream and are being fixed.
• MoE AutoEP (permute_rows / unit / fused_parity, 228): 225 pass / 3 skip (the DeepEP-transport tests are CUDA-only). The async split-plan parity test whose guard this branch relaxes (accelerator instead of CUDA) was additionally verified in a multi-card NPU run.
• test_chunked_cross_entropy.py (50): 47 pass / 3 skip (the skips are CUDA-only memory-peak, fork-safety, and device-assert subprocess tests). The Triton backend is now exercised on the accelerator instead of only CUDA.
• Total: 564 pass / 11 fail / 8 skip — the 11 failures are the upstream triton-ascend issues noted above.

The torch.device(...) launch-context replacement is exercised by the NPU run of the packed SwiGLU and fused RoPE paths.

Hardware: Ascend 910B4 (single device for the suites; the async split-plan parity test additionally ran multi-card), torch 2.10.0+cpu, torch_npu 2.10.0.post7, CANN 8.5.1, triton-ascend 3.6.0.

Pre-commit

pre-commit run --files <the 12 changed files> — all checks passed (yapf, flake8, check-torchdist, check-torchcuda, check-license, codespell).

[chatgpt-codex-connector]

Codex Review Summary

This comment shows the latest Codex review activity on this pull request.

Review	Status	Commit	Review trigger
📝 Code Review	⚠️ Failed 2026-10-20T13:42:27.850252Z	5cc7afe	Draft marked ready

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review" or "@codex security review".

Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings.

[Copilot]

🟡 Changes recommended

CPU eligibility regressions and the missing advertised permutation-dispatch update remain unresolved.

3 open issues

• Reject CPU tensors in CPU_Accelerator.on_accelerator · New
• Exclude CPU from CPU_Accelerator.is_available · New
• Enable Triton MoE reorder dispatch for supported accelerators · New

What changed in this PR

Generalizes DeepSpeed’s Triton guards and device contexts to support non-CUDA accelerators with an active Triton backend.

Changes:

• Adds shared accelerator-tensor checks and current-device switching.
• Enables NPU Triton detection and broadens runtime eligibility.
• Adapts kernel and AutoEP tests for multiple accelerator types.

File	Description
tests/​unit/​v1/​runtime/​test_chunked_cross_entropy.py	Broadens backend coverage and adapts rejection checks.
tests/​unit/​v1/​ops/​triton_ops/​test_swiglu_triton.py	Generalizes availability and device-context tests.
tests/​unit/​v1/​ops/​triton_ops/​test_group_gemm_triton.py	Uses accelerator capability checks.
tests/​unit/​v1/​ops/​triton_ops/​test_fused_rotary_pos_emb.py	Generalizes device tests and layout comparisons.
tests/​unit/​v1/​ops/​triton_ops/​test_fused_rms_norm.py	Updates availability, errors, and layout checks.
tests/​unit/​v1/​ops/​triton_ops/​test_autoep_fused_token_ops.py	Broadens eligibility and updates error assertions.
tests/​unit/​v1/​moe/​test_autoep_unit.py	Generalizes device cases and protocol doubles.
tests/​unit/​v1/​moe/​test_autoep_permute_rows.py	Enables non-CUDA accelerator cases.
tests/​unit/​v1/​moe/​test_autoep_grad_parity.py	Broadens async split-plan test eligibility.
deepspeed/​runtime/​chunked_cross_entropy.py	Uses accelerator-aware backend selection.
deepspeed/​ops/​triton_ops/​swiglu_triton.py	Uses shared current-device contexts.
deepspeed/​ops/​triton_ops/​fused_rotary_pos_emb.py	Generalizes tensor guards and launch contexts.
deepspeed/​ops/​triton_ops/​fused_rms_norm.py	Uses the shared tensor guard.
deepspeed/​ops/​triton_ops/​autoep_fused_token_ops.py	Generalizes guards and device errors.
deepspeed/​ops/​triton_ops/​_triton.py	Adds tensor eligibility and device-context helpers.
deepspeed/​module_inject/​auto_ep_layer.py	Broadens async split-plan tensor eligibility.
accelerator/​npu_accelerator.py	Delegates Triton support to the shared probe.

_{Review effort: Balanced}

[jinyouzhi]

Very significant!
Just curious, what do you think about the discussion in #8789?