[AI][Ascend] Guard native ACLNN calls with the tensor device

[yswang777]

Issue Type

• Bug Report
• Feature Request
• Operator Implementation
• Platform Support
• Performance Issue
• Documentation

AI Agent Information

• Agent: OpenAI Codex
• Model: GPT-5
• Session Context: Multi-device Qwen-Image-2512 execution on Ascend required explicit torch.flagos.set_device() calls at transformer-shard and scheduler boundaries to prevent NaN outputs.

Summary

Ascend native ACLNN calls obtain their stream and workspace device from the ambient current device rather than the device of their tensor arguments. After a multi-device module finishes on its last shard, a subsequent operation on a tensor owned by another flagos:N device can launch or allocate against the stale device context. The Qwen-Image scheduler produced NaNs until every shard boundary and scheduler update manually selected the input tensor's device.

Environment

Click to expand environment details

• Platform: Ascend
• Python: 3.11.15
• PyTorch: 2.10.0+cpu
• torch_fl: fe8fe3d5d6674ef357fae03810d11c9a59635b2d (bug observed); current main checked at 90a184f06d5792d9cbdeeaf08e19e8eb813fe574
• Hardware: Ascend 910C, multiple visible logical devices

Reproduction

import torch_fl
import torch

torch.flagos.set_device(1)
x = torch.tensor([1.0, 2.0, 4.0], dtype=torch.float32, device="flagos:0")

# Representative native/FlagGems scalar work performed after another shard
# left device 1 current. The final regression should select a confirmed ACLNN
# native op from the Ascend configuration and assert both values and context.
y = x + 0.5
torch.flagos.synchronize(0)

print("result", y.cpu().tolist())
print("current", torch.flagos.current_device())

The Qwen reproducer leaves the last transformer shard current and then calls FlowMatchEulerDiscreteScheduler.step() on latents owned by the scheduler/VAE device. Without explicit device selection the run produced NaN/black output; selecting the latent device before scheduler operations made all 12 Showcase cases finite.

Expected vs Actual Behavior

Expected: an operator executes on the device of its input/output tensors, and a temporary device guard restores the caller's previous current device after the call.

Actual: ACLNN execution uses GetCurrentAclStream() and workspace allocation uses the ambient current device. Multi-device callers currently need Python hooks that call torch.flagos.set_device() around normal tensor operations.

Root Cause Analysis

Generated and handwritten Ascend kernels allocate outputs from tensor options, but EXEC_ASCEND_CMD obtains the launch stream through GetCurrentAclStream(), which reads ::GetDevice. ExecAscendCached also keys executors and allocates workspace from ::GetDevice. Neither layer derives or guards the current device from the first tensor argument. Output allocation can temporarily guard its requested device and then restore the stale ambient device before the ACLNN launch, leaving the launch and workspace on the wrong logical device.

Open PR #322 fixes an analogous current-device leak for CUDA-compatible boxing guards. It does not cover Ascend native ACLNN kernels, so this issue is related but not duplicate.

Proposed Solution

Add one conditional device guard at a shared Ascend-native dispatch boundary, based on the first defined tensor argument. It must be active for both generated/cached and handwritten ACLNN paths, restore the previous device, handle factory/copy operations separately, and avoid a runtime set-device call when the tensor device already matches the current device.

The exact insertion point needs a focused prototype: dispatcher-level guarding covers every native kernel, while generator/macro-level guarding gives more control over cross-device operators. The fix must not blindly force all tensor-list or copy operators onto the first input device.

Verification Plan

• First add a deterministic two-device operator reproducer that fails against the unfixed build.
• Cover generated cached kernels and at least one handwritten EXEC_ASCEND_CMD kernel.
• Assert result values, result device, stream/device selection, and restoration of the caller's current device.
• Cover same-device calls and benchmark guard overhead.
• Audit copy and mixed-device operators so the guard does not change their contract.
• Remove the Qwen Python context hooks and rerun the complete Showcase.

Context & Investigation

1. Preserved Qwen runs that produced NaN/black images when the first latent/scheduler operation followed transformer sharding without selecting its device.
2. Added narrow Python hooks at shard, transformer-return, scheduler, and initial-latent boundaries; all 12 cases then completed with finite latents.
3. Read empty.cc, the PrivateUse1 guard implementation, ExecAscendCached, EXEC_ASCEND_CMD, and the Ascend stream registry.
4. Confirmed the launch helpers use ambient ::GetDevice rather than tensor device.
5. Checked open work: PR fix: keep the flagos current device across boxed calls #322 concerns CUDA-compatible boxing and does not touch Ascend ACLNN.

Related Code Locations

• csrc/aten/backends/ascend/op_api_common.h:493 - cached executor reads ambient device.
• csrc/aten/backends/ascend/op_api_common.h:662 - EXEC_ASCEND_CMD uses GetCurrentAclStream().
• csrc/runtime/accelerator/ascend/stream_api.cc:130 - current stream resolves from ::GetDevice.
• csrc/aten/empty.cc:31 - output allocation conditionally guards a requested device.
• csrc/runtime/guard.h:75 - PrivateUse1 device guard implementation.

Checklist - AI Agents MUST Complete All

• I have provided complete environment information
• I have included a minimal reproducer and the real workload signature
• I have included the observed failure mode
• I have analyzed the root cause
• I have proposed a specific bounded design
• I have identified affected code locations
• I have described how to verify the fix
• I have checked for duplicate issues and PRs
• All text is in English
• Code follows project conventions