# Copyright (c) PyPTO Contributors.
# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
# CANN Open Software License Agreement Version 2.0 (the "License").
# Please refer to the License for details. You may not use this file except in compliance with the License.
# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
# See LICENSE in the root of the software repository for the full text of the License.
# -----------------------------------------------------------------------------------------------------------
"""Opt-in NPU isolation for DSpark's persistent TP communication protocol.
Run under task-submit with four or sixteen devices in TASK_DEVICE.
"""
import os
import sys
import time
from pathlib import Path
import pytest
def _program(payload, padded, single_writer):
import pypto.language as pl
import pypto.language.distributed as pld
tp = 4
rows = 128
width = 4096
slots = 24 if single_writer else 1
stride = 32 if padded else 1
signal_width = slots * stride
world = pl.dynamic("WORLD")
@pl.jit
def chip(
inputs: pl.Tensor[[rows, width], pl.BF16],
output: pl.Out[pl.Tensor[[tp * rows, width], pl.BF16]],
counts: pl.Out[pl.Tensor[[2, tp * slots], pl.INT32]],
skew: pl.Tensor[[1], pl.INT32],
probe: pl.Out[pl.Tensor[[1], pl.INT32]],
window: pld.DistributedTensor[[tp * rows, width], pl.BF16],
hidden_done: pld.DistributedTensor[[tp, signal_width], pl.INT32],
logits_done: pld.DistributedTensor[[tp, signal_width], pl.INT32],
base: pl.Scalar[pl.INT32],
local: pl.Scalar[pl.INT32],
):
with pl.at(level=pl.Level.CORE_GROUP, name_hint="rank_skew") as delay_tid:
iterations = pl.read(skew, [0])
skew_value = iterations
for index in pl.range(iterations):
skew_value = pl.cast((skew_value * 17 + index) % 65521, pl.INT32)
pl.write(probe, [0], skew_value)
with pl.spmd(8, name_hint="hidden_push", deps=[delay_tid]) as push_tid:
block = pl.tile.get_block_idx()
if payload:
for peer in pl.range(tp):
pld.tensor.put(
dst=window, peer=base + peer, src=inputs,
dst_offsets=[local * rows + block * 16, 0],
src_offsets=[block * 16, 0], shape=[16, width],
)
for peer in pl.range(tp):
if peer != local:
if single_writer:
pld.system.notify(
target=hidden_done, peer=base + peer,
offsets=[local, block * stride], value=1, op=pld.NotifyOp.Set,
)
else:
pld.system.notify(
target=hidden_done, peer=base + peer,
offsets=[local, 0], value=1, op=pld.NotifyOp.AtomicAdd,
)
with pl.at(level=pl.Level.CORE_GROUP, name_hint="hidden_wait", deps=[push_tid]) as wait_tid:
for peer in pl.range(tp):
if peer != local:
if single_writer:
for block in pl.range(8):
pld.system.wait(
signal=hidden_done, offsets=[peer, block * stride],
expected=1, cmp=pld.WaitCmp.Ge,
)
else:
pld.system.wait(
signal=hidden_done, offsets=[peer, 0], expected=8, cmp=pld.WaitCmp.Ge,
)
for peer in pl.range(tp):
for slot in pl.range(slots):
hidden_value = pl.read(hidden_done, [peer, slot * stride])
pl.write(counts, [0, peer * slots + slot], hidden_value)
with pl.spmd(32, name_hint="hidden_gather", deps=[wait_tid, push_tid]) as gather_tid:
block = pl.tile.get_block_idx()
r0 = block // 4 * 64
c0 = block % 4 * 1024
if payload:
output[r0:r0 + 64, c0:c0 + 1024] = window[r0:r0 + 64, c0:c0 + 1024]
else:
output[r0:r0 + 64, c0:c0 + 1024] = pl.full([64, 1024], dtype=pl.BF16, value=0.0)
with pl.spmd(24, name_hint="logits_notify", deps=[gather_tid]) as notify_tid:
block = pl.tile.get_block_idx()
for peer in pl.range(tp):
if peer != local:
if single_writer:
pld.system.notify(
target=logits_done, peer=base + peer,
offsets=[local, block * stride], value=1, op=pld.NotifyOp.Set,
)
else:
pld.system.notify(
target=logits_done, peer=base + peer,
offsets=[local, 0], value=1, op=pld.NotifyOp.AtomicAdd,
)
with pl.at(level=pl.Level.CORE_GROUP, name_hint="logits_wait", deps=[notify_tid]) as finish_tid:
for peer in pl.range(tp):
if peer != local:
if single_writer:
for slot in pl.range(slots):
pld.system.wait(
signal=logits_done, offsets=[peer, slot * stride],
expected=1, cmp=pld.WaitCmp.Ge,
)
else:
pld.system.wait(
signal=logits_done, offsets=[peer, 0], expected=24, cmp=pld.WaitCmp.Ge,
)
for peer in pl.range(tp):
for slot in pl.range(slots):
logits_value = pl.read(logits_done, [peer, slot * stride])
pl.write(counts, [1, peer * slots + slot], logits_value)
with pl.at(level=pl.Level.CORE_GROUP, name_hint="signal_clear", deps=[finish_tid]):
for peer in pl.range(tp):
for slot in pl.range(slots):
pl.write(hidden_done, [peer, slot * stride], pl.cast(0, pl.INT32))
pl.write(logits_done, [peer, slot * stride], pl.cast(0, pl.INT32))
@pl.jit.host
def distributed(
inputs: pl.Tensor[[world, rows, width], pl.BF16],
output: pl.Out[pl.Tensor[[world, tp * rows, width], pl.BF16]],
counts: pl.Out[pl.Tensor[[world, 2, tp * slots], pl.INT32]],
skew: pl.Tensor[[world, 1], pl.INT32],
probe: pl.Out[pl.Tensor[[world, 1], pl.INT32]],
):
window_buf = pld.alloc_window_buffer(tp * rows * width * 2)
hidden_buf = pld.alloc_window_buffer(tp * signal_width * 4)
logits_buf = pld.alloc_window_buffer(tp * signal_width * 4)
for rank in pl.range(pld.world_size()):
window = pld.window(window_buf, [tp * rows, width], dtype=pl.BF16)
hidden_done = pld.window(hidden_buf, [tp, signal_width], dtype=pl.INT32)
logits_done = pld.window(logits_buf, [tp, signal_width], dtype=pl.INT32)
chip(
inputs[rank], output[rank], counts[rank], skew[rank], probe[rank],
window, hidden_done, logits_done, rank // tp * tp, rank % tp, device=rank,
)
return distributed
def _split_phase_program():
import pypto.language as pl
import pypto.language.distributed as pld
world = pl.dynamic("WORLD")
@pl.jit
def chip(
phase: pl.Tensor[[1], pl.INT32],
counts: pl.Out[pl.Tensor[[2, 4], pl.INT32]],
hidden_done: pld.DistributedTensor[[4, 1], pl.INT32],
logits_done: pld.DistributedTensor[[4, 1], pl.INT32],
base: pl.Scalar[pl.INT32],
local: pl.Scalar[pl.INT32],
):
with pl.spmd(8, name_hint="split_phase"):
block = pl.tile.get_block_idx()
step = pl.read(phase, [0])
if step == 0:
for peer in pl.range(4):
if peer != local:
pld.system.notify(
target=hidden_done, peer=base + peer, offsets=[local, 0],
value=1, op=pld.NotifyOp.AtomicAdd,
)
for repeat in pl.range(3):
pld.system.notify(
target=logits_done, peer=base + peer, offsets=[local, 0],
value=1, op=pld.NotifyOp.AtomicAdd,
)
else:
if block == 0:
for peer in pl.range(4):
hidden = pl.read(hidden_done, [peer, 0])
logits = pl.read(logits_done, [peer, 0])
pl.write(counts, [0, peer], hidden)
pl.write(counts, [1, peer], logits)
pl.write(hidden_done, [peer, 0], pl.cast(0, pl.INT32))
pl.write(logits_done, [peer, 0], pl.cast(0, pl.INT32))
@pl.jit.host
def distributed(
phase: pl.Tensor[[world, 1], pl.INT32],
counts: pl.Out[pl.Tensor[[world, 2, 4], pl.INT32]],
):
hidden_buf = pld.alloc_window_buffer(16)
logits_buf = pld.alloc_window_buffer(16)
for rank in pl.range(pld.world_size()):
hidden = pld.window(hidden_buf, [4, 1], dtype=pl.INT32)
logits = pld.window(logits_buf, [4, 1], dtype=pl.INT32)
chip(phase[rank], counts[rank], hidden, logits, rank // 4 * 4, rank % 4, device=rank)
return distributed
def test_split_phase_atomic():
if not os.environ.get("TASK_DEVICE"):
pytest.skip("requires task-submit device allocation")
import torch
from pypto.ir import DistributedConfig
from pypto.runtime import RunConfig
devices = [int(item) for item in os.environ["TASK_DEVICE"].split(",")]
assert len(devices) in (4, 16), devices
torch.set_num_threads(2)
ranks = len(devices)
root = Path("run-logs/comm-isolation/split_phase")
root.mkdir(parents=True, exist_ok=True)
config = RunConfig(
platform="a2a3", save_kernels=True, save_kernels_dir=str(root / "kernels"),
distributed_config=DistributedConfig(device_ids=devices, num_sub_workers=0, aicpu_thread_num=4),
ring_heap=256 * 1024 * 1024,
)
phase = torch.zeros((ranks, 1), dtype=torch.int32).share_memory_()
counts = torch.zeros((ranks, 2, 4), dtype=torch.int32).share_memory_()
expected = torch.tensor([8, 24], dtype=torch.int32).reshape(1, 2, 1).repeat(ranks, 1, 4)
for rank in range(ranks):
expected[rank, :, rank % 4] = 0
compiled = _split_phase_program().compile(phase, counts, config=config)
with compiled.prepare(persistent=True, reset_persistent_windows=False) as worker:
for iteration in range(100):
phase.zero_()
worker(phase, counts)
phase.fill_(1)
counts.fill_(-1)
worker(phase, counts)
assert torch.equal(counts, expected), f"split phase round={iteration} counters={counts.tolist()}"
if iteration % 10 == 0 or iteration == 99:
print(f"ISOLATION split_phase round={iteration + 1}/100 PASS", flush=True)
@pytest.mark.parametrize("payload,padded,single_writer", [
pytest.param(False, False, False, id="atomic_compact"),
pytest.param(False, True, False, id="atomic_padded"),
pytest.param(False, True, True, id="atomic_set"),
pytest.param(True, False, False, id="payload_compact"),
pytest.param(True, True, False, id="payload_padded"),
pytest.param(True, True, True, id="payload_set"),
])
def test_persistent_communication(payload, padded, single_writer):
if not os.environ.get("TASK_DEVICE"):
pytest.skip("requires task-submit device allocation")
import torch
from pypto.ir import DistributedConfig
from pypto.runtime import RunConfig
devices = [int(item) for item in os.environ["TASK_DEVICE"].split(",")]
assert len(devices) in (4, 16), devices
torch.set_num_threads(2)
ranks = len(devices)
slots = 24 if single_writer else 1
variant = f"payload{int(payload)}_pad{int(padded)}_set{int(single_writer)}"
root = Path("run-logs/comm-isolation") / variant
root.mkdir(parents=True, exist_ok=True)
config = RunConfig(
platform="a2a3", save_kernels=True, save_kernels_dir=str(root / "kernels"),
distributed_config=DistributedConfig(device_ids=devices, num_sub_workers=0, aicpu_thread_num=4),
ring_heap=256 * 1024 * 1024,
)
inputs = torch.zeros((ranks, 128, 4096), dtype=torch.bfloat16).share_memory_()
output = torch.zeros((ranks, 512, 4096), dtype=torch.bfloat16).share_memory_()
counts = torch.zeros((ranks, 2, 4 * slots), dtype=torch.int32).share_memory_()
skew = torch.zeros((ranks, 1), dtype=torch.int32).share_memory_()
probe = torch.zeros((ranks, 1), dtype=torch.int32).share_memory_()
args = (inputs, output, counts, skew, probe)
program = _program(payload, padded, single_writer)
compiled = program.compile(*args, config=config)
expected_counts = torch.zeros_like(counts)
for rank in range(ranks):
for peer in range(4):
if peer != rank % 4:
if single_writer:
expected_counts[rank, 0, peer * slots:peer * slots + 8] = 1
expected_counts[rank, 1, peer * slots:(peer + 1) * slots] = 1
else:
expected_counts[rank, :, peer] = torch.tensor([8, 24])
pattern = torch.arange(128 * 4096, dtype=torch.int32).reshape(128, 4096) % 97
started = time.monotonic()
with compiled.prepare(persistent=True, reset_persistent_windows=False) as worker:
for iteration in range(100):
for rank in range(ranks):
inputs[rank].copy_(((pattern + rank * 7 + iteration * 13) % 127).to(torch.bfloat16))
skew[rank, 0] = 200000 if iteration % 2 and rank % 4 == (iteration // 2) % 4 else 0
counts.fill_(-1)
output.fill_(-1)
worker(*args)
assert torch.equal(counts, expected_counts), (
f"{variant} round={iteration} counters={counts.tolist()} expected={expected_counts.tolist()}"
)
for rank in range(ranks):
if payload:
base = rank // 4 * 4
expected = inputs[base:base + 4].reshape(512, 4096)
assert torch.equal(output[rank], expected), f"{variant} round={iteration} rank={rank} payload"
else:
assert torch.count_nonzero(output[rank]).item() == 0
if iteration % 10 == 0 or iteration == 99:
print(f"ISOLATION {variant} round={iteration + 1}/100 PASS elapsed={time.monotonic() - started:.2f}s", flush=True)
if __name__ == "__main__":
raise SystemExit(pytest.main([__file__, "-v", "-s", *sys.argv[1:]]))
Platform
a2a3 (Ascend 910B/C hardware)
Runtime Variant
tensormap_and_ringbuffer
Description
A model-free, 16-rank persistent TP notification test hung on its first invocation on ranks 0-3, while ranks 4-15 completed. About 404.5 seconds later, the affected ranks reported ACL 507901 (
hdc disconnect). Domain cleanup then failed on all 16 ranks, and the process remained alive after pytest printed its failure summary.This is an intermittent observed failure, not a deterministic reproducer or a confirmed AtomicAdd defect. Subsequent compact-counter runs and all communication controls passed. The failing run used a locally instrumented Simpler checkout; the complete instrumentation diff is attached below. Hardware/driver failure, instrumentation effects, and runtime scheduling remain unresolved possibilities.
This investigation started from DSpark serving hangs, but the isolated failure's signature is different: the original serving log had scheduler S1 stalls in a Markov graph; this isolated run did not report S1 or TENSOR_WAIT_TIMEOUT. A common root cause has not been established.
Steps to Reproduce
Prerequisites: Linux aarch64, 16 allocated A3 devices, Python 3.12 with torch/torch_npu and pytest, the pinned PyPTO/Simpler/PTO ISA versions below, and PTOAS 0.60. No model checkpoint, serving installation, or pypto-lib is needed for the attached test. PyPTO must be configured to use this Simpler installation and PTO ISA checkout.
Place the attached Python source at
tests/integration/test_dspark_comm_isolation.pyin a writable working directory. Apply the attached diagnostic patch to the pinned Simpler checkout and install it using that checkout's normal build procedure to match the tested environment. An unmodified-runtime comparison is still needed; the patch is supplied for reproducibility, not proposed as a fix.Run the compact case with the timeout settings from the failed attempt:
The task must provide comma-separated physical device IDs through
TASK_DEVICE. SetPTO_ISA_ROOTto the pinned PTO ISA checkout in the job environment if it is not already configured. The test supports four devices too, but all results reported here used 16 devices, forming four independent TP=4 groups.Run the complete control matrix by removing
-k atomic_compactand-x. For shorter failure diagnosis, the successful controls usedSIMPLER_OP_EXECUTE_TIMEOUT_US=60000000,SIMPLER_STREAM_SYNC_TIMEOUT_MS=90000, andSIMPLER_SCHEDULER_TIMEOUT_MS=20000.The compact case runs 100 host-fenced iterations with persistent windows and
reset_persistent_windows=False. Within each TP group, eight producer blocks AtomicAdd 1 to each peer's source-rank counter, peers TWAIT >=8, a dependent gather runs, then 24 blocks AtomicAdd 1 to the second counter and peers TWAIT >=24. Both sets of counters are cleared locally before the invocation returns. Each compact signal array contains four adjacent INT32 counters. Alternate iterations deliberately delay a rotating rank in each TP group.Controls independently change counter layout (128-byte stride), replace shared AtomicAdd counters with per-producer Set slots, add 4 MiB of exact-checked cross-rank BF16 payload traffic, or separate notification and counter read/reset into host-fenced invocations with no device TWAIT. All output/counter assertions run on every iteration. The atomic-only case has no remote payload transfer.
Expected Behavior
Every rank completes; remote counter totals are exactly 8 and 24, and payload controls match the rank/iteration-specific source data exactly. After an execution failure, cleanup should terminate without leaving the process alive or obscuring the primary failure with a domain-release exception.
Actual Behavior
Failed attempt: first
atomic_compactround, no completed-round PASS marker. Ranks 4-15 reportednode.complete ... run_id=1 ... outcome=0. Ranks 0-3 reported failed completion only after approximately 404.5 seconds. The configured device operation timeout was 400 seconds; the elapsed time is compatible with that limit, but does not locate where execution stopped.Representative host diagnostics:
The process still held the allocation after the pytest summary and was terminated through the task scheduler at about 10 minutes. There were no counter readbacks on the failed invocation because execution failure skipped copy-back. No device-side TWAIT counter values or exact stuck PC were captured, so these logs cannot establish a lost notification.
Successful controls, all on 16 devices:
Observed compact-case frequency: one failed initialization/run out of eight attempts; the seven successful attempts completed 700 rounds total. Six successful attempts used the shorter timeouts and one used the original long timeouts. This is not a controlled estimate of the underlying failure probability.
The Set controls initially encountered a test compile error caused by reusing a scalar variable name across scopes; the attached source uses distinct scalar names, and both controls were then compiled and passed. Those compile failures are not counted as runtime failures.
Git Commit ID
Simpler:
4e4d3a4ad1e54c1db3d50e72decc025a9075bfa0, plus the attached uncommitted scheduler instrumentation patch.PyPTO:
747e1e4ba99f13dea1a609d893ac7512e85a5897.PTO ISA:
a8040450238f162985d8b596fbebeb54bfba2bf5.PTOAS:
0.60.CANN Version
9.0.0
Driver Version
26.0.rc1
Host Platform
Linux (aarch64)
Additional Context
Related: #1798, #2136, #824. These are related scheduling/communication investigations, not confirmed duplicates.
Source inspection found explicit DCCI/barrier operations in the current TNOTIFY/TWAIT path and generated payload publication/clear code. Passing controls do not prove the full serving graph is correct, but they provide no positive evidence of a general AtomicAdd, compact-counter false-sharing, or payload visibility defect. Please help determine whether the observed hang originates in runtime dispatch, communication, the diagnostic register reads, or the device/driver path.
The runtime patch adds WAITLIST diagnostics and live DATA_MAIN_BASE/PMU register reads in stall formatting. It changes observability and may affect timing; its relevance to this failure has not been isolated.
Complete standalone test: test_dspark_comm_isolation.py
Exact local Simpler diagnostic patch used for these runs