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

A backend-level gate is, however, too coarse once a backend is open: a kernel-level backend bug would still crash at runtime. The PR therefore adds a second, per-kernel gate: each Triton call site validates its own kernel on first use through a runtime probe — a minimal forward+backward run on the live device, cached per process — so the platform answers for itself kernel by kernel, no blacklist to maintain, and a backend fix flips the gate automatically. Finally, where the fix is not a gate answer but an implementation difference, the PR adds the third piece, per-hardware kernel variants through an annotation on the public function: every DeepSpeed-owned Triton kernel entry on the training path carries a dispatch decorator and is itself the generic fallback; an accelerator names its own triton_ops subpackage (triton_ops_dir, op_builder-style) and a platform module overrides the function by name — the first variant, shipped in this PR, lowers the Ascend packed SwiGLU backward to BLOCK_SIZE=1024 (the workaround for the triton-ascend single-element unified-buffer compile failure). CUDA remains the primary, fully supported platform, and every default path behaves exactly as before.

Summary

Layer 1 — device-agnostic guards and the shared driver probe. 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.

Layer 2 — per-kernel runtime probes. The Triton gate becomes per-kernel and self-reporting:

• is_triton_supported(kernel=None, probe=None) gains optional per-kernel parameters on CUDA and NPU. Without arguments it keeps today's backend-level meaning; with a probe, the backend actually compiles and runs the kernel's minimal risky usage on the live device once per process and caches the answer, keyed by kernel:device:dtype.
• Every Triton call site on the training path — swiglu, swiglu_packed, group_gemm, fused RMSNorm, fused RoPE, fused weighted_restore / row_weighting, chunked cross entropy, fp8 GEMM, liger vocab-parallel CE, and permute_rows — routes through this gate with its probe defined in a single library (deepspeed/ops/triton_ops/_probes.py), so probes never scatter across call sites.
• A platform whose gate predates the extension keeps its static answer; probe failures are fail-closed with a one-time warning; DS_TRITON_KERNEL_GATE=0 is an emergency escape hatch back to the backend-level gate.
• The probe directly catches the known NPU issues: the fp8 GEMM gate closes automatically on NPU (no float8 tensor support) instead of crashing at runtime, and the permute_rows gather path remains a narrow, probe-unsafe static declaration in the per-kernel gate (the NPU accelerator answers False for the kernel name; multi-card device fault, upstream issue in the comment).

Layer 3 — per-hardware kernel variants, dispatched by annotation. Where the fix is an implementation difference rather than a gate answer, the platform ships its own variant:

• The portable public function carries @dispatch("<feature>") and is itself the generic fallback; a platform overrides it by defining a same-named function in its triton_ops subpackage (re-exporting the decorated function counts as "no override", so a variant can never recurse). Every DeepSpeed-owned Triton kernel entry on the training path carries the annotation — swiglu / swiglu_packed, grouped GEMM, fused RMSNorm / RoPE, fused weighted-restore / row-weighting, vocab-parallel softmax, chunked cross entropy, fp8 GEMM, and the MoE permute / gather kernels — so a platform variant plugs in without touching any call site. (The liger vocab-parallel CE wrapper is excluded: it wraps a third-party kernel, and the gate already covers it.) Kernels outside triton_ops pass their portable module to the decorator; platform variants still live under the accelerator's triton_ops subpackage.
• abstract_accelerator.triton_ops_dir() answers the accelerator's triton_ops subpackage (None = portable, the default); the NPU accelerator points at deepspeed.ops.triton_ops.npu. The resolver loads the platform module with a portable fallback and caches per feature; a subpackage may hold a single kernel.
• The steps no platform should fork — validation, the runtime probe gate, the eager fallback — are private helpers of the portable module that the variant imports, so a variant only forks what actually differs.
• First variant (this PR): the Ascend packed SwiGLU backward lowers its block size to 1024 — the workaround for the triton-ascend single-element unified-buffer compile failure, numerically verified on a 910B4 (forward and the separate kernels stay at 2048). The kernel probes resolve the class through the same dispatch as the production path, so the gate tests exactly the implementation the production path runs.
• The single-element packed-SwiGLU special case is removed: with the Ascend variant the shape compiles, so the shared code keeps no shape-specific gate. Platforms without a variant see no 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.

A backend-level gate cannot catch kernel-level backend bugs

5. Opening the NPU gate opens every kernel at once. Once triton-ascend is installed, every Triton path runs on NPU. But the NPU suite found kernel-level failures a boolean "backend installed" gate cannot express: the packed SwiGLU backward fails to compile for one packed element (MLIRCompilationError, upstream unified-buffer overflow), several fp32/fp16 kernels differ from eager by 1 ulp, and the permute_rows gather faults the device in multi-card runs. A backend-level gate either leaves the whole platform closed (the old return False) or opens every kernel including the broken ones.
6. Three call sites are gated only backend-level (fp8_gemm, liger vocab-parallel CE, the inference attention flag), so they flip open on NPU without any per-kernel validation.

Fix

Layer 1 — device-agnostic guards

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

Layer 2 — per-kernel runtime probes

• accelerator/abstract_accelerator.py: is_triton_supported() keeps its contract.
• cuda_accelerator.py / npu_accelerator.py: is_triton_supported(kernel=None, probe=None) — the backend-level answer without arguments; with a probe, the shared probe_triton_kernel helper compiles, synchronizes, and caches. npu_accelerator.py declares the gather exception for the permute_rows_gather kernel name (return False, upstream issue in the comment).
• _triton.py: probe_triton_kernel(key, run, reference=None, tolerance=None) — one run per process; a current-stream synchronization after the run (not a device-wide one) so asynchronous faults fail the probe here, not on the first real workload, without waiting on unrelated streams or surfacing their errors; optional bitwise / allclose comparison against a reference; thread-safe cache; fail-closed with a one-time warning. is_triton_kernel_supported(kernel, dtype, device, probe_fn, *args) builds the probe only on a cache miss, so the steady-state hot path is one dict lookup, defers the probe while the current stream is being captured (CUDA graphs), and supports a probe-less static declaration (probe_fn=None) where the accelerator answers for the kernel name alone — the probe-unsafe paths.
• _probes.py (new): one probe per kernel — a minimal forward+backward on the caller's device and dtype, wrapped in the device context. Probe inputs are deterministic constants, not random draws, so a probe never touches the caller's RNG state by construction and is reproducible run over run. (Necessary: on NPU, torch_npu 2.10.0.post7's set_rng_state is a no-op and device Generator objects share the default generator's state, so no fork/save-restore-based isolation can be sound there.)
• Call sites: swiglu / swiglu_packed fall back to eager; group_gemm_triton falls back to torch._grouped_mm; fused RMSNorm / RoPE / weighted-restore / row-weighting raise their existing "unsupported" errors, which their dispatchers already turn into the eager path; chunked CE resolves auto to torch when the probe fails; fp8_gemm and the liger gate take the same per-kernel probe.
• permute_rows_supported routes through the per-kernel gate as a probe-less static declaration (is_triton_kernel_supported(..., None)) instead of device.type == "cuda".
• Tests: a new tests/unit/v1/ops/triton_ops/test_triton_probe.py pins the gate contract (run-once cache, fail-closed cached answer, RNG isolation, escape hatch) on any accelerator with an active Triton backend.

Layer 3 — per-hardware kernel variants, dispatched by annotation

• accelerator/abstract_accelerator.py: triton_ops_dir() returns None by default — backends opt in only when they ship a variant, so the 11 existing platforms need no change.
• accelerator/npu_accelerator.py: triton_ops_dir() returns deepspeed.ops.triton_ops.npu.
• _triton.py: resolve_triton_ops_impl(feature, portable_module=None) — asks the accelerator, imports the platform module, falls back to the portable one on a miss, and caches per feature — and the dispatch(feature, portable_module=None) decorator: the decorated portable function runs the platform's same-named function when it exists, the portable original otherwise (steady state: one dict lookup + one getattr per call).
• triton_ops/npu/swiglu_triton.py (new): defines swiglu_packed — validation, the runtime gate, and the eager fallback are the portable module's shared helpers, and only the packed backward's BLOCK_SIZE is lowered to 1024 (portable kernel source reused unchanged); the module stays importable without a device or a Triton install, so resolution never depends on hardware.
• swiglu_triton.py: swiglu and swiglu_packed carry @dispatch("swiglu_triton"); _validate_packed_input / _swiglu_packed_eager / _swiglu_packed_supported are factored out for the variant; the _swiglu_packed_1elem_supported special case and its probe are deleted (the shared code keeps no shape-specific gate).
• fused_rms_norm.py / fused_rotary_pos_emb.py / autoep_fused_token_ops.py / group_gemm_triton.py / vocab_parallel_softmax.py: every public kernel entry carries @dispatch("<module>").
• deepspeed/runtime/chunked_cross_entropy.py, deepspeed/ops/fp_quantizer/fp8_gemm_triton.py, deepspeed/moe/ep_kernels.py: entries carry @dispatch("<feature>", portable_module=<own module>); the resolver's portable path takes the module name, so out-of-package kernels dispatch the same way with their variants under the platform's triton_ops subpackage.
• _probes.py: the packed probe resolves the class through the same dispatch (it must not call the public function itself — that would re-enter the gate), so the gate tests the platform's variant.
• Tests: tests/unit/v1/ops/triton_ops/test_triton_dispatch.py (new) pins the resolver contract with a duck-typed accelerator double — no subpackage resolves portable, a subpackage resolves the platform module, an unoverridden feature falls back to portable, a missing feature raises, an out-of-package feature resolves its explicit portable module — plus three decorator contract tests — the platform override runs, no override falls back to the original, a re-exported wrapper counts as "no override" — on any host, no device required.

Runtime-detection mechanics

First use of a kernel	the probe compiles and runs the minimal case once; the answer is cached per kernel:device:dtype
Every later call	one dict lookup; the probe is neither run nor constructed
Probe failure	fallback path + one-time warning; DS_TRITON_KERNEL_GATE=0 disables probing globally
Upstream fix lands	the gate flips automatically on the next process, zero code change
New platform	answers for itself on first use; nothing to register anywhere

Known limitations (deliberate, documented)

• Probes are canaries: they validate that the backend can compile and run this kernel family, not the numerics at production shapes — production-shape correctness stays the job of the unit-test suites, which this PR runs on both platforms.
• The probe catches Python-raised failures (compile errors such as the Ascend MLIRCompilationError are Python exceptions). Hard native crashes (segfault/abort) are outside the try/except model; none has been observed, and subprocess isolation is a possible follow-up if one appears.
• Failure policy is intentionally asymmetric: explicitly opted-in config paths (combine_impl="fused_weighted_sum", row_weighting_impl="fused") fail fast with a clear error, while drop-in APIs (swiglu, group_gemm, chunked CE, fp8) fall back silently — each matches its caller's contract.
• The permute_rows gather is the single static per-kernel declaration: probing a device fault is unsafe, so the NPU accelerator's is_triton_supported answers False for the kernel name until the upstream fix lands.
• The Ascend packed-SwiGLU variant is the single per-hardware kernel variant; it exists only to work around the upstream unified-buffer failure and is removed when the fix lands.

Follow-ups (explicitly out of scope)

• Subprocess isolation with a timeout for probes, if a native crash or hang ever appears (the current probes catch Python-raised failures, which is all that has been observed).
• The failure warning could advise that the cached answer lives for the process lifetime (restart to re-probe); DS_TRITON_KERNEL_GATE could gain a force-enable mode; transient-failure retries are not evaluated.
• An optional warmup entry point for eager pre-probing, and per-tensor-device capability checks for the CUDA static gate in mixed-device setups.
• Remove the Ascend packed-SwiGLU variant once triton-ascend fixes the single-element unified-buffer overflow (or keep it if 1024 stays the better Ascend tile). Every entry point already carries the dispatch annotation, so platforms can bring further variants without touching shared code; the flash-linear-attention triton-ascend playbook (fixed launch configs on Ascend instead of autotune, smaller backward tiles for the unified buffer) is the reference template for them.
• The torch_npu RNG-state findings above (set_rng_state no-op, non-independent device Generators) are worth filing upstream with torch_npu.

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. The only added runtime cost is one minimal kernel compile per kernel family per process (measured below).
• 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 (318): 302 pass / 9 fail / 7 skip — 8 of the 9 failures are the triton-ascend backend issues reported upstream and being fixed (6 packed-SwiGLU 1-ulp bitwise-vs-separate, the exact baseline set, + 2 row-weighting 1-ulp eager-parity with the bit-identical 2.29e-05 / 3.05e-05 deltas); the 9th is an environmental torch.OutOfMemoryError in test_expert_offset_exceeds_int32 — the test allocates ~4.3 GiB, the device currently has ~1.3 GiB free (94% HBM occupied by colocated workloads), and the test is untouched by this PR. The 7 skips are 4 two-device cases (single-card run), 2 CUDA-only cases, and 1 memory-guarded parity case. The 13 new in-repo gate/dispatch contract tests (5 probe + 8 dispatch) are included and pass (the 4 tests over the previous run are the new dispatch decorator contract tests). The single-element packed-SwiGLU shapes now run the packed path through the Ascend variant (BLOCK_SIZE=1024) instead of any fallback.
• MoE AutoEP (permute_rows / unit / fused_parity / grad_parity, 259): 242 pass / 17 skip / 0 fail — the skips are the multi-device and distributed cases (full-step and grad-parity scenarios need >1 device; the run is single-card, matching the earlier runs). 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: 591 pass / 9 fail / 27 skip — 8 of the 9 failures are the upstream triton-ascend issues noted above and the 9th is the environmental OOM noted above; this commit introduces 0 new failures (the 13 gate/dispatch contract tests pass, and the swiglu failure set is exactly the baseline).

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.

Probe layer (NPU, Ascend 910B4, triton-ascend 3.6.0)

Before/after the probe layer on the affected four files:

	before	after
packed SwiGLU 1-element × 3 dtypes (upstream compile crash)	3 fail (MLIRCompilationError)	3 pass (probe closes the gate → separate-kernel fallback)
packed SwiGLU fp32/fp16 bitwise-vs-separate (upstream 1-ulp semantics)	6 fail	6 fail — identical to before
everything else	350 pass	350 pass
total	9 fail / 350 pass	6 fail / 353 pass

Full training-path sweep (10 files, 583 tests): 565 pass / 18 fail, where the 18 are exactly
6 swiglu + 2 row-weighting upstream-numerics failures (the row-weighting deltas 2.29e-05 /
3.05e-05 are bit-identical at the pre-probe baseline commit) plus 10 artifacts of editing files
mid-run; rerunning the affected files cleanly gives 8 pre-existing failures, 0 introduced.

Probe cost (NPU; each measurement in a fresh subprocess with an empty Triton compile cache;
the benchmark script is kept out of the repo):

measurement	result
marginal cost of the first-use probe (vs. the kernel's own first compile, ~30 s on triton-ascend)	0–2.4 s per kernel, once per process
steady-state gate cost per call (cached)	1.7–1.9 µs (one dict lookup)
end-to-end gate overhead on swiglu 32×16 bf16	~9 µs/call (~5% of that small kernel; sub-1% at MoE shapes)
RNG isolation	probe inputs are deterministic constants — the probe draws no random numbers, so the caller's generator is never touched (pinned by the in-repo control test)

Two probe catches verified end-to-end: the packed SwiGLU 1-element probe absorbs the
MLIRCompilationError and the three tests go green via the fallback; the fp8 probe closes
the gate on NPU (no float8_e4m3fn support) where the open gate previously crashed at runtime.
Both still hold after the probes switched to deterministic inputs.

Per-hardware dispatch (NPU, Ascend 910B4, triton-ascend 3.6.0)

• resolve_triton_ops_impl("swiglu_triton") resolves deepspeed.ops.triton_ops.npu.swiglu_triton on NPU, and the variant overrides swiglu_packed by name — the decorated portable function dispatches to the variant's own function and keeps the portable implementation as its fallback, with _PACKED_BWD_BLOCK_SIZE == 1024 (pinned on-device by the verification script below).
• The single-element packed SwiGLU runs end to end through the public entry point into the dispatched variant (the portable 2048 tile does not compile for this shape): forward and backward match eager within tolerance (allclose; packed-vs-eager is not bitwise on Ascend, the upstream numerics issue).
• The swiglu suite keeps the exact baseline failure set — 6 known 1-ulp bitwise failures, 55 pass, 2 skip — with the single-element shapes now running the packed path instead of a fallback.
• test_triton_dispatch.py (8): all pass, no device required — also verified in a Triton-less process, which additionally proves every decorated module imports cleanly without a Triton install and the dispatched chains answer through their eager fallbacks.
• test_triton_probe.py (5): all pass (re-run after the dispatch change).
• The on-device script also pins all nine feature resolutions: the Ascend variant for swiglu_triton, the portable module for the other eight.

Verification script (kept out of the repo): test_npu/deepspeed/roadmap2026q4/triton_ascend/verify_dispatch_npu.py — ALL CHECKS PASSED on Ascend 910B4.

Probe-gate unit tests (in-repo, platform-agnostic)

tests/unit/v1/ops/triton_ops/test_triton_probe.py pins the gate's contract — the
wall-clock costs are the benchmark numbers above, not CI assertions:

• test_probe_runs_once_and_later_calls_answer_from_the_cache — the probe factory runs exactly once; later calls answer from the cache without running or constructing the probe.
• test_failing_probe_closes_the_gate_with_a_cached_false — a raising probe closes the gate with a cached False, and later calls keep the fallback without re-running any probe.
• test_probe_does_not_advance_the_caller_rng_sequence — control comparison: the first draw after seeding, with the probe run in between, must equal the same draw without the probe. On NPU this test is red against the previous fork-based probes and green with the deterministic inputs (verified on Ascend 910B4).
• test_env_escape_hatch_bypasses_probing — DS_TRITON_KERNEL_GATE=0 answers at the backend level and never constructs a probe.
• test_static_declaration_answers_from_the_backend_gate_without_a_probe — a probe-less static declaration (probe_fn=None) runs nothing and answers with the backend-level gate for a kernel name no backend declares.

Result: 5 pass / 0 fail on Ascend 910B4 (triton-ascend 3.6.0), 5.6 s.

Pre-commit

pre-commit run --files <the 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-23T07:13:30.444567Z We couldn’t complete this review. Comment @codex review to try again.	a2a9888	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?

[VenusTZZ]

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

Hi, @jinyouzhi .Sorry I hadn't noticed issue #8789 earlier — it's a great discussion and raises some really good points.

I was wondering if we could improve is_triton_supported() by having generic probes check Triton compatibility on the running hardware at runtime, so each platform doesn't have to keep its own blacklist. I've tried this out in the current PR and also left a note over in #8789.

Feel free to take a look when you get a chance — happy to rework it if this approach doesn't fit.

[jinyouzhi]

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

Hi, @jinyouzhi .Sorry I hadn't noticed issue #8789 earlier — it's a great discussion and raises some really good points.

I was wondering if we could improve is_triton_supported() by having generic probes check Triton compatibility on the running hardware at runtime, so each platform doesn't have to keep its own blacklist. I've tried this out in the current PR and also left a note over in #8789.

Feel free to take a look when you get a chance — happy to rework it if this approach doesn't fit.

Thank you for your reply. I think this idea and practice very useful and potential to be genetic for all backends. Could you also share your opinion and results into the issue?

[VenusTZZ]

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

Hi, @jinyouzhi .Sorry I hadn't noticed issue #8789 earlier — it's a great discussion and raises some really good points.
I was wondering if we could improve is_triton_supported() by having generic probes check Triton compatibility on the running hardware at runtime, so each platform doesn't have to keep its own blacklist. I've tried this out in the current PR and also left a note over in #8789.
Feel free to take a look when you get a chance — happy to rework it if this approach doesn't fit.

Thank you for your reply. I think this idea and practice very useful and potential to be genetic for all backends. Could you also share your opinion and results into the issue?

Of course, very happy to discuss this further! I've already posted the details in #8789