Skip to content

[Bug] sync_start drain can deadlock when generation-less ack/election state is reused across retry attempts #1455

Closed
#1461
Closed
[Bug] sync_start drain can deadlock when generation-less ack/election state is reused across retry attempts#1455
#1461
Labels
bugSomething isn't working

Description

@Leaf-Salix
Leaf-Salix
opened on Jul 24, 2026

Platform

a2a3 (Ascend 910B/C hardware)

Runtime Variant

tensormap_and_ringbuffer

Description

sync_start drain appears to use reusable shared coordination fields without an attempt/generation tag:

sync_start_pending
drain_ack_mask
drain_worker_elected
drain_stage_go
drain_stage_done_mask

When the elected drain owner observes insufficient resources, it resets drain_ack_mask and drain_worker_elected so scheduler threads can resume completion polling and retry the same pending sync_start task.

However, a scheduler thread from the old attempt can already have passed the ack barrier and then be delayed before election. After the owner resets the shared fields, that stale thread can continue and interpret the reset drain_worker_elected == 0 / drain_ack_mask values as if they belonged to its own attempt.

This is a generation-less reusable-barrier ABA issue: every individual shared field value is valid, but the fields can come from different drain attempts.

I used a deterministic test-only harness to force this interleaving. The harness does not directly return a synthetic error code. It lets the stale thread continue into the original handle_drain_mode() election/stage protocol. The original protocol then enters a split barrier and eventually fails with an AICPU timeout.

One detail in the harness is intentionally deterministic: after proving that the stale thread crossed from attempt 1 to attempt 2, it lets the stale owner pass the availability check. This models the real timing where resources become available before the stale owner performs that check. The hook does not replace the ack/election/stage protocol; the deadlock still comes from the original protocol accepting a stale owner and publishing stage_go with a mixed-generation ack_mask.

Steps to Reproduce

1. Apply the env-gated runtime hook patch from Additional Context. The patch adds `SIMPLER_DRAIN_ABA_TEST`, tags drain attempts in test mode, and pauses one scheduler thread after the old attempt's ack barrier but before election.

2. Add the ST from Additional Context. The orchestration submits one holder MIX task and then one `sync_start` MIX task:


static constexpr int16_t HOLDER_BLOCKS = 1;
static constexpr int16_t SYNC_BLOCKS = 24;
static constexpr int64_t HOLDER_SPIN_ITERS = 50000000;

submit_mix(holder, scratch, HOLDER_BLOCKS, 0, HOLDER_SPIN_ITERS, false);
submit_mix(output, scratch, SYNC_BLOCKS, 0, 0, true);


The test config used:


{
    "aicpu_thread_num": 4,
    "block_dim": 24,
    "runtime_env": {
        "ring_heap": 32 * 1024 * 1024,
        "ring_task_window": 1024,
    },
}


3. Run the test on a2a3 hardware:


SIMPLER_DRAIN_ABA_TEST=1 pytest \
  tests/st/a2a3/tensormap_and_ringbuffer/spmd_sync_start_drain_aba_hook/test_spmd_sync_start_drain_aba_hook.py \
  -s -v --forked --platform=a2a3

Expected Behavior

The drain retry protocol should not mix scheduler threads from different drain attempts.

A stale scheduler thread from an old drain attempt should not be able to:

  • win election for a newer attempt,
  • publish stage_go using a mixed-generation ack_mask,
  • or wait for stage_done from threads that are actually waiting in the next attempt's ack barrier.

Expected behavior would be one of:

  • stale-attempt threads detect a generation mismatch and return to completion polling/retry;
  • or the drain protocol tags ack/election/stage state with a generation id so fields from different attempts cannot be combined.

Actual Behavior

The test fails by AICPU timeout:

RuntimeError: run failed with code 507018

Key logs from the failing run:

DRAIN_ABA_DIAG retry_reset ... thread=2 task=1 shape=2 gated=0 global_available=23 block_num=24 ack_mask=0x7 elected=3 stage_go=0 stage_done=0x0

DRAIN_ABA_TEST releasing stale drain attempt into original protocol:
  thread=1 task=1 local_attempt=1 current_attempt=2 block_num=24 ack_mask=0x5 elected=0

DRAIN_ABA_TEST forcing stale owner past availability check:
  thread=1 task=1 observed_available=23 block_num=24 ack_mask=0x5 elected=2

DRAIN_ABA_DIAG stage_go ... thread=1 task=1 shape=2 gated=0 global_available=24 block_num=24 ack_mask=0x5

DRAIN_ABA_DIAG stage_done ... thread=1 task=1 gated=0 my_running=21 ... done_before=0x0

DRAIN_ABA_DIAG ack_wait ... thread=2 task=1 ack_mask=0x5 all_acked=0x7 elected=2 stage_go=1 sync_start_pending=24
DRAIN_ABA_DIAG ack_wait ... thread=0 task=1 ack_mask=0x5 all_acked=0x7 elected=2 stage_go=1 sync_start_pending=24

The final state is split across two attempts:

thread1: stale owner from the old attempt, waiting for stage_done == 0x7
thread0/thread2: new attempt ack barrier, waiting for ack_mask == 0x7
actual ack_mask: 0x5

Both waits are individually valid for their own attempt, but they are now waiting on different generations. This causes a deadlock.

Git Commit ID

d0bc661

CANN Version

Ascend CANN toolkit 9.0.0

Driver Version

26.0.rc1

Host Platform

Linux (aarch64)

Additional Context

The test-only harness is only used to make the interleaving deterministic. It is not the failure mechanism itself:

  • the harness first proves the stale thread crossed from local_attempt=1 to current_attempt=2;
  • then it releases that stale thread into the original handle_drain_mode() protocol;
  • the original protocol publishes stage_go with ack_mask=0x5;
  • two peer scheduler threads remain in the new attempt's ack barrier, waiting for ack_mask=0x7.

A likely fix is to add an explicit drain generation/attempt tag. Each scheduler thread should record the generation it observed before acking. Election, availability check, stage_go, stage-done, and finalize should all verify the generation still matches. On insufficient-resource retry, the owner should advance the generation before clearing/reusing ack_mask and drain_worker_elected, so old-attempt threads can safely return instead of entering the next attempt's election/stage protocol.

Tested and suspected commits

The failing run was tested against upstream hw-native-sys/simpler commit d0bc661a5b41b4925f2d72c77e86228cc351cf23. I have not completed a full bisection, so the following is only static-history analysis:

Complete runtime hook/instrumentation patch

This is the full runtime hook and instrumentation patch used for the deterministic run. The important interleaving hook is the SIMPLER_DRAIN_ABA_TEST path in handle_drain_mode(); the additional DRAIN_ABA_DIAG logs were used to confirm the final split wait.

Show full runtime hook patch 
diff --git a/src/a2a3/runtime/tensormap_and_ringbuffer/host/runtime_maker.cpp b/src/a2a3/runtime/tensormap_and_ringbuffer/host/runtime_maker.cpp
index c3f00927..0308a833 100644
--- a/src/a2a3/runtime/tensormap_and_ringbuffer/host/runtime_maker.cpp
+++ b/src/a2a3/runtime/tensormap_and_ringbuffer/host/runtime_maker.cpp
@@ -608,9 +608,13 @@ static void apply_orch_sched_env_flags(Runtime *runtime) {
     const char *serial_env = std::getenv("SIMPLER_TMR_SERIAL_ORCH_SCHED_ENABLE");
     runtime->dev.serial_orch_sched =
         serial_env && (serial_env[0] == '1' || serial_env[0] == 't' || serial_env[0] == 'T');
+    const char *drain_aba_env = std::getenv("SIMPLER_DRAIN_ABA_TEST");
+    runtime->dev.drain_aba_test_mode =
+        drain_aba_env && (drain_aba_env[0] == '1' || drain_aba_env[0] == 't' || drain_aba_env[0] == 'T');
     LOG_INFO_V0(
         "Serial orchestrator-to-scheduler start gate: %s", runtime->dev.serial_orch_sched ? "enabled" : "disabled"
     );
+    LOG_INFO_V0("Drain ABA test hook: %s", runtime->dev.drain_aba_test_mode ? "enabled" : "disabled");
 }
 
 // per-(cid,config): reserve and acquire the static device pools. GM heap, PTO2
diff --git a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_ring_buffer.h b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_ring_buffer.h
index 070074ce..fbbaca10 100644
--- a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_ring_buffer.h
+++ b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_ring_buffer.h
@@ -453,6 +453,84 @@ private:
                 static_cast<int>(h.task_state.load(std::memory_order_acquire)), rc & ~PTO2_FANOUT_SCOPE_BIT,
                 fc & ~PTO2_FANOUT_SCOPE_BIT, (rc & PTO2_FANOUT_SCOPE_BIT) ? 1 : 0
             );
+            int32_t dump_end = std::min(local_task_id_, last_alive + 9);
+            for (int32_t tid = last_alive; tid < dump_end; tid++) {
+                PTO2TaskSlotState &s = slot_states_[tid & window_mask_];
+                PTO2TaskPayload *p = s.payload;
+                uint32_t s_fc = s.fanout_count;
+                uint32_t s_rc = s.fanout_refcount.load(std::memory_order_acquire);
+                int32_t staged_cores = 0;
+                int32_t published = -1;
+                int32_t early_dispatch = -1;
+                int32_t launch = -1;
+                int32_t drain = -1;
+                int32_t running_slot = -1;
+                if (p != nullptr) {
+                    for (int w = 0; w < PTO2_EARLY_DISPATCH_CORE_MASK_WORDS; w++) {
+                        staged_cores += __builtin_popcountll(p->staged_core_mask[w].load(std::memory_order_seq_cst));
+                    }
+                    published = static_cast<int32_t>(p->published_block_count.load(std::memory_order_seq_cst));
+                    early_dispatch = static_cast<int32_t>(p->early_dispatch_state.load(std::memory_order_seq_cst));
+                    launch = static_cast<int32_t>(p->early_dispatch_launch_state.load(std::memory_order_seq_cst));
+                    drain = static_cast<int32_t>(p->early_sync_drain_state.load(std::memory_order_seq_cst));
+                    running_slot = static_cast<int32_t>(p->running_slot_count.load(std::memory_order_seq_cst));
+                }
+                LOG_ERROR(
+                    "  DRAIN_ABA_DIAG slot local=%d raw=%lld state=%d fanin=%d/%d fanout=%u/%u scope=%d "
+                    "subtasks=%d/%d next_block=%d logical=%d published=%d staged_cores=%d running_slot=%d "
+                    "early_dispatch=%d launch=%d drain=%d",
+                    tid, s.task != nullptr ? static_cast<long long>(s.task->task_id.raw) : -1LL,
+                    static_cast<int>(s.task_state.load(std::memory_order_acquire)),
+                    s.fanin_refcount.load(std::memory_order_acquire), s.fanin_count, s_rc & ~PTO2_FANOUT_SCOPE_BIT,
+                    s_fc & ~PTO2_FANOUT_SCOPE_BIT, (s_rc & PTO2_FANOUT_SCOPE_BIT) ? 1 : 0,
+                    static_cast<int32_t>(s.completed_subtasks.load(std::memory_order_acquire)),
+                    static_cast<int32_t>(s.total_required_subtasks),
+                    static_cast<int32_t>(s.next_block_idx.load(std::memory_order_seq_cst)),
+                    static_cast<int32_t>(s.logical_block_num), published, staged_cores, running_slot, early_dispatch,
+                    launch, drain
+                );
+                if (p != nullptr) {
+                    int32_t inline_count = std::min(p->fanin_actual_count, PTO2_FANIN_INLINE_CAP);
+                    for (int32_t fi = 0; fi < inline_count; fi++) {
+                        PTO2TaskSlotState *prod = p->fanin_inline_slot_states[fi];
+                        uint32_t prod_fc = prod != nullptr ? prod->fanout_count : 0;
+                        uint32_t prod_rc =
+                            prod != nullptr ? prod->fanout_refcount.load(std::memory_order_acquire) : 0;
+                        LOG_ERROR(
+                            "  DRAIN_ABA_DIAG slot local=%d fanin[%d]=raw:%lld state:%d fanout:%u/%u scope:%d",
+                            tid, fi,
+                            prod != nullptr && prod->task != nullptr ? static_cast<long long>(prod->task->task_id.raw) :
+                                                                        -1LL,
+                            prod != nullptr ? static_cast<int>(prod->task_state.load(std::memory_order_acquire)) : -1,
+                            prod_rc & ~PTO2_FANOUT_SCOPE_BIT, prod_fc & ~PTO2_FANOUT_SCOPE_BIT,
+                            (prod_rc & PTO2_FANOUT_SCOPE_BIT) ? 1 : 0
+                        );
+                    }
+                }
+                int32_t fanout_dump_count = 0;
+                for (PTO2DepListEntry *edge = s.fanout_head; edge != nullptr && fanout_dump_count < 8;
+                     edge = edge->next, fanout_dump_count++) {
+                    PTO2TaskSlotState *consumer = edge->slot_state;
+                    int32_t consumer_local =
+                        consumer != nullptr ? static_cast<int32_t>(consumer - slot_states_) : -1;
+                    LOG_ERROR(
+                        "  DRAIN_ABA_DIAG slot local=%d fanout[%d]=local:%d raw:%lld state:%d fanin:%d/%d "
+                        "subtasks:%d/%d",
+                        tid, fanout_dump_count, consumer_local,
+                        consumer != nullptr && consumer->task != nullptr ?
+                            static_cast<long long>(consumer->task->task_id.raw) :
+                            -1LL,
+                        consumer != nullptr ? static_cast<int>(consumer->task_state.load(std::memory_order_acquire)) :
+                                              -1,
+                        consumer != nullptr ? consumer->fanin_refcount.load(std::memory_order_acquire) : -1,
+                        consumer != nullptr ? consumer->fanin_count : -1,
+                        consumer != nullptr ?
+                            static_cast<int32_t>(consumer->completed_subtasks.load(std::memory_order_acquire)) :
+                            -1,
+                        consumer != nullptr ? static_cast<int32_t>(consumer->total_required_subtasks) : -1
+                    );
+                }
+            }
         }
         LOG_ERROR("Solution:");
         if (scope_gated) {
diff --git a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/runtime.h b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/runtime.h
index 90977d3e..427a0923 100644
--- a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/runtime.h
+++ b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/runtime.h
@@ -187,6 +187,12 @@ struct alignas(64) DeviceRuntimeLaunchDesc {
     // Controlled via SIMPLER_TMR_SERIAL_ORCH_SCHED_ENABLE environment variable.
     bool serial_orch_sched;
 
+    // Test-only drain interleaving hook. When enabled, scheduler threads create a
+    // deterministic stale-attempt window in sync_start drain to prove whether the
+    // reusable ack/election fields are generation-less.
+    // Controlled via SIMPLER_DRAIN_ABA_TEST environment variable.
+    bool drain_aba_test_mode;
+
     void *gm_sm_ptr_;                        // GM pointer to PTO2 shared memory (device)
     ChipStorageTaskArgs orch_args_storage_;  // Copy of args for device
 
diff --git a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_cold_path.cpp b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_cold_path.cpp
index 7d4e1a82..9e4f989e 100644
--- a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_cold_path.cpp
+++ b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_cold_path.cpp
@@ -1130,6 +1130,7 @@ int32_t SchedulerContext::pre_handshake_init(
     aic_count_ = cores_total_num_ / 3;
     aiv_count_ = (cores_total_num_ * 2) / 3;
     active_sched_threads_ = (sched_thread_num_ > 0) ? sched_thread_num_ : aicpu_thread_num_;
+    drain_aba_test_mode_ = runtime->dev.drain_aba_test_mode;
     handshake_failed_.store(false, std::memory_order_release);
 
     // State the barrier-free per-thread init path no longer reaches via
@@ -1311,6 +1312,9 @@ void SchedulerContext::deinit() {
     drain_state_.drain_stage_go.store(0, std::memory_order_release);
     drain_state_.drain_stage_done_mask.store(0, std::memory_order_release);
     drain_state_.drain_running_staged.store(0, std::memory_order_release);
+    drain_state_.drain_attempt_seq.store(0, std::memory_order_release);
+    drain_state_.drain_test_victim_armed.store(0, std::memory_order_release);
+    drain_state_.drain_test_stale_seen.store(0, std::memory_order_release);
     drain_state_.pending_task.store(nullptr, std::memory_order_release);
 
     // Reset task counters and orchestrator state
@@ -1324,6 +1328,7 @@ void SchedulerContext::deinit() {
     aiv_count_ = 0;
     cores_total_num_ = 0;
     aicpu_thread_num_ = 0;
+    drain_aba_test_mode_ = false;
     sched_thread_num_ = 0;
     active_sched_threads_ = 0;
     for (int32_t t = 0; t < MAX_AICPU_THREADS; t++) {
diff --git a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_completion.cpp b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_completion.cpp
index 1b162a61..ac2ae387 100644
--- a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_completion.cpp
+++ b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_completion.cpp
@@ -18,6 +18,7 @@
 #include "common/l2_swimlane_profiling.h"
 #include "common/memory_barrier.h"
 #include "common/platform_config.h"
+#include "pto_runtime_status.h"
 #include "pto_runtime2.h"
 #include "runtime.h"
 #include "spin_hint.h"
@@ -240,6 +241,23 @@ void SchedulerContext::complete_slot_task(
             }
             deferred_release_slot_states[deferred_release_count++] = &slot_state;
         }
+        if (slot_state.task_attrs.requires_sync_start()) {
+            static std::atomic<int32_t> drain_aba_sync_release_logs{0};
+            int32_t log_idx = drain_aba_sync_release_logs.fetch_add(1, std::memory_order_relaxed);
+            if (log_idx < 256) {
+                LOG_INFO_V9(
+                    "DRAIN_ABA_DIAG sync_complete_release sample=%d thread=%d task=%lld count=%d",
+                    log_idx, thread_idx, static_cast<long long>(slot_state.task->task_id.raw), deferred_release_count
+                );
+            }
+            while (deferred_release_count > 0) {
+#if SIMPLER_SCHED_PROFILING
+                (void)sched_->on_task_release(*deferred_release_slot_states[--deferred_release_count], thread_idx);
+#else
+                sched_->on_task_release(*deferred_release_slot_states[--deferred_release_count]);
+#endif
+            }
+        }
         completed_this_turn++;
     }
 
@@ -514,6 +532,11 @@ bool SchedulerContext::enter_drain_mode(PTO2TaskSlotState *slot_state, int32_t b
     drain_state_.drain_stage_go.store(0, std::memory_order_relaxed);
     drain_state_.drain_stage_done_mask.store(0, std::memory_order_relaxed);
     drain_state_.drain_running_staged.store(0, std::memory_order_relaxed);
+    if (drain_aba_test_mode_) {
+        drain_state_.drain_attempt_seq.store(1, std::memory_order_release);
+        drain_state_.drain_test_victim_armed.store(0, std::memory_order_relaxed);
+        drain_state_.drain_test_stale_seen.store(0, std::memory_order_relaxed);
+    }
     // Release store: all stores above are now visible to any thread that
     // acquire-loads sync_start_pending and sees block_num > 0.
     drain_state_.sync_start_pending.store(block_num, std::memory_order_release);
@@ -702,21 +725,117 @@ void SchedulerContext::handle_drain_mode(int32_t thread_idx, [[maybe_unused]] ui
         block_num = drain_state_.sync_start_pending.load(std::memory_order_acquire);
     } while (block_num < 0);
     if (block_num == 0) return;
+    const int32_t local_drain_attempt =
+        drain_aba_test_mode_ ? drain_state_.drain_attempt_seq.load(std::memory_order_acquire) : 0;
 
     uint32_t all_acked = (1u << active_sched_threads_) - 1;
+    {
+        static std::atomic<int32_t> drain_aba_handle_enter_logs{0};
+        int32_t log_idx = drain_aba_handle_enter_logs.fetch_add(1, std::memory_order_relaxed);
+        PTO2TaskSlotState *pending = drain_state_.pending_task.load(std::memory_order_acquire);
+        bool later_task = pending != nullptr && pending->task != nullptr && pending->task->task_id.local() != 2;
+        if (log_idx < 256 || later_task) {
+            LOG_INFO_V9(
+                "DRAIN_ABA_DIAG handle_enter sample=%d thread=%d task=%lld block_num=%d ack_mask=0x%x "
+                "elected=%d stage_go=%d stage_done=0x%x sync_start_pending=%d",
+                log_idx, thread_idx,
+                pending != nullptr && pending->task != nullptr ? static_cast<long long>(pending->task->task_id.raw) :
+                                                                  -1LL,
+                block_num, drain_state_.drain_ack_mask.load(std::memory_order_relaxed),
+                drain_state_.drain_worker_elected.load(std::memory_order_relaxed),
+                drain_state_.drain_stage_go.load(std::memory_order_relaxed),
+                drain_state_.drain_stage_done_mask.load(std::memory_order_relaxed),
+                drain_state_.sync_start_pending.load(std::memory_order_relaxed)
+            );
+        }
+    }
 
     // Ack barrier -- signal this thread has stopped dispatch.
-    drain_state_.drain_ack_mask.fetch_or(1u << thread_idx, std::memory_order_release);
+    uint32_t ack_after = drain_state_.drain_ack_mask.fetch_or(1u << thread_idx, std::memory_order_release) |
+                         (1u << thread_idx);
+    {
+        static std::atomic<int32_t> drain_aba_ack_logs{0};
+        int32_t log_idx = drain_aba_ack_logs.fetch_add(1, std::memory_order_relaxed);
+        PTO2TaskSlotState *pending = drain_state_.pending_task.load(std::memory_order_acquire);
+        bool later_task = pending != nullptr && pending->task != nullptr && pending->task->task_id.local() != 2;
+        if (log_idx < 256 || later_task) {
+            LOG_INFO_V9(
+                "DRAIN_ABA_DIAG ack_after sample=%d thread=%d task=%lld ack_mask=0x%x all_acked=0x%x "
+                "elected=%d stage_go=%d",
+                log_idx, thread_idx,
+                pending != nullptr && pending->task != nullptr ? static_cast<long long>(pending->task->task_id.raw) :
+                                                                  -1LL,
+                ack_after, all_acked, drain_state_.drain_worker_elected.load(std::memory_order_relaxed),
+                drain_state_.drain_stage_go.load(std::memory_order_relaxed)
+            );
+        }
+    }
 
     // Spin until all threads have acked.
     // If our bit is cleared while waiting, elected reset due to insufficient resources.
+    int32_t ack_wait_spins = 0;
     while (true) {
         if (is_completed()) return;
         uint32_t ack = drain_state_.drain_ack_mask.load(std::memory_order_acquire);
         if ((ack & all_acked) == all_acked) break;
         if ((ack & (1u << thread_idx)) == 0) return;
+        ack_wait_spins++;
+        if (ack_wait_spins == (1 << 20)) {
+            static std::atomic<int32_t> drain_aba_ack_wait_logs{0};
+            int32_t log_idx = drain_aba_ack_wait_logs.fetch_add(1, std::memory_order_relaxed);
+            PTO2TaskSlotState *pending = drain_state_.pending_task.load(std::memory_order_acquire);
+            bool later_task = pending != nullptr && pending->task != nullptr && pending->task->task_id.local() != 2;
+            if (log_idx < 128 || later_task) {
+                LOG_INFO_V9(
+                    "DRAIN_ABA_DIAG ack_wait sample=%d thread=%d task=%lld ack_mask=0x%x all_acked=0x%x "
+                    "elected=%d stage_go=%d sync_start_pending=%d",
+                    log_idx, thread_idx,
+                    pending != nullptr && pending->task != nullptr ?
+                        static_cast<long long>(pending->task->task_id.raw) :
+                        -1LL,
+                    ack, all_acked, drain_state_.drain_worker_elected.load(std::memory_order_relaxed),
+                    drain_state_.drain_stage_go.load(std::memory_order_relaxed),
+                    drain_state_.sync_start_pending.load(std::memory_order_relaxed)
+                );
+            }
+        }
         SPIN_WAIT_HINT();
     }
+
+    if (drain_aba_test_mode_ && active_sched_threads_ >= 3 && thread_idx == 1) {
+        int32_t expected_armed = 0;
+        if (drain_state_.drain_test_victim_armed.compare_exchange_strong(
+                expected_armed, 1, std::memory_order_acq_rel, std::memory_order_relaxed
+            )) {
+            int32_t spin_count = 0;
+            int32_t current_attempt = drain_state_.drain_attempt_seq.load(std::memory_order_acquire);
+            while (current_attempt == local_drain_attempt && !is_completed() && spin_count < (1 << 24)) {
+                ++spin_count;
+                SPIN_WAIT_HINT();
+                current_attempt = drain_state_.drain_attempt_seq.load(std::memory_order_acquire);
+            }
+            if (current_attempt != local_drain_attempt) {
+                PTO2TaskSlotState *pending = drain_state_.pending_task.load(std::memory_order_acquire);
+                uint32_t peer_acked = all_acked & ~(1u << thread_idx);
+                uint32_t ack = drain_state_.drain_ack_mask.load(std::memory_order_acquire);
+                while (!is_completed() && ((ack & peer_acked) != peer_acked || (ack & (1u << thread_idx)) != 0)) {
+                    SPIN_WAIT_HINT();
+                    ack = drain_state_.drain_ack_mask.load(std::memory_order_acquire);
+                }
+                drain_state_.drain_test_stale_seen.store(1, std::memory_order_release);
+                LOG_ERROR(
+                    "DRAIN_ABA_TEST releasing stale drain attempt into original protocol: thread=%d task=%lld "
+                    "local_attempt=%d current_attempt=%d block_num=%d ack_mask=0x%x elected=%d spin_count=%d",
+                    thread_idx,
+                    pending != nullptr && pending->task != nullptr ? static_cast<long long>(pending->task->task_id.raw) :
+                                                                     -1LL,
+                    local_drain_attempt, current_attempt, block_num, ack,
+                    drain_state_.drain_worker_elected.load(std::memory_order_relaxed), spin_count
+                );
+            }
+        }
+    }
+
     // Election -- exactly one thread wins the CAS.
     int32_t expected = 0;
     drain_state_.drain_worker_elected.compare_exchange_strong(
@@ -744,12 +863,51 @@ void SchedulerContext::handle_drain_mode(int32_t thread_idx, [[maybe_unused]] ui
         // needs block_num idle cores/clusters.
         int32_t available =
             count_global_available(shape, slot_state->active_mask.core_mask(), /*include_pending=*/gated);
+        if (drain_aba_test_mode_ && thread_idx == 1 &&
+            drain_state_.drain_test_stale_seen.load(std::memory_order_acquire) == 1) {
+            LOG_ERROR(
+                "DRAIN_ABA_TEST forcing stale owner past availability check: thread=%d task=%lld observed_available=%d "
+                "block_num=%d ack_mask=0x%x elected=%d",
+                thread_idx, static_cast<long long>(slot_state->task->task_id.raw), available, block_num,
+                drain_state_.drain_ack_mask.load(std::memory_order_relaxed),
+                drain_state_.drain_worker_elected.load(std::memory_order_relaxed)
+            );
+            available = block_num;
+        }
         if (available < block_num) {
+            static std::atomic<int32_t> drain_aba_retry_logs{0};
+            int32_t log_idx = drain_aba_retry_logs.fetch_add(1, std::memory_order_relaxed);
+            if (log_idx < 256 || slot_state->task->task_id.local() != 2) {
+                LOG_INFO_V9(
+                    "DRAIN_ABA_DIAG retry_reset sample=%d thread=%d task=%lld shape=%d gated=%d global_available=%d "
+                    "block_num=%d ack_mask=0x%x elected=%d stage_go=%d stage_done=0x%x",
+                    log_idx, thread_idx, static_cast<long long>(slot_state->task->task_id.raw),
+                    static_cast<int32_t>(shape), static_cast<int32_t>(gated), available, block_num,
+                    drain_state_.drain_ack_mask.load(std::memory_order_relaxed),
+                    drain_state_.drain_worker_elected.load(std::memory_order_relaxed),
+                    drain_state_.drain_stage_go.load(std::memory_order_relaxed),
+                    drain_state_.drain_stage_done_mask.load(std::memory_order_relaxed)
+                );
+            }
             // Insufficient -- reset so all threads resume completion polling to free cores, then retry.
+            if (drain_aba_test_mode_) {
+                drain_state_.drain_attempt_seq.fetch_add(1, std::memory_order_acq_rel);
+            }
             drain_state_.drain_ack_mask.store(0, std::memory_order_release);
             drain_state_.drain_worker_elected.store(0, std::memory_order_release);
             return;
         }
+        static std::atomic<int32_t> drain_aba_stage_logs{0};
+        int32_t log_idx = drain_aba_stage_logs.fetch_add(1, std::memory_order_relaxed);
+        if (log_idx < 128) {
+            LOG_INFO_V9(
+                "DRAIN_ABA_DIAG stage_go sample=%d thread=%d task=%lld shape=%d gated=%d global_available=%d "
+                "block_num=%d ack_mask=0x%x",
+                log_idx, thread_idx, static_cast<long long>(slot_state->task->task_id.raw), static_cast<int32_t>(shape),
+                static_cast<int32_t>(gated), available, block_num,
+                drain_state_.drain_ack_mask.load(std::memory_order_relaxed)
+            );
+        }
         // Release parallel staging: every thread (this one included) now stages its own cores.
         drain_state_.drain_running_staged.store(0, std::memory_order_relaxed);
         drain_state_.drain_stage_done_mask.store(0, std::memory_order_relaxed);
@@ -772,6 +930,31 @@ void SchedulerContext::handle_drain_mode(int32_t thread_idx, [[maybe_unused]] ui
     if (drain_prof) drain_acked_ts = get_sys_cnt_aicpu();  // pre-stage
 #endif
     int32_t my_running = drain_stage_cores(slot_state, block_num, thread_idx, gated);
+    {
+        static std::atomic<int32_t> drain_aba_stage_done_logs{0};
+        int32_t log_idx = drain_aba_stage_done_logs.fetch_add(1, std::memory_order_relaxed);
+        if (log_idx < 96) {
+            int32_t staged_cores = 0;
+            if (slot_state->payload != nullptr) {
+                for (int w = 0; w < PTO2_EARLY_DISPATCH_CORE_MASK_WORDS; w++) {
+                    staged_cores += __builtin_popcountll(
+                        slot_state->payload->staged_core_mask[w].load(std::memory_order_seq_cst)
+                    );
+                }
+            }
+            LOG_INFO_V9(
+                "DRAIN_ABA_DIAG stage_done sample=%d thread=%d task=%lld gated=%d my_running=%d next_block=%d "
+                "published=%d staged_cores=%d done_before=0x%x running_staged_before=%d",
+                log_idx, thread_idx, static_cast<long long>(slot_state->task->task_id.raw), static_cast<int32_t>(gated),
+                my_running, slot_state->next_block_idx.load(std::memory_order_seq_cst),
+                slot_state->payload != nullptr ?
+                    static_cast<int32_t>(slot_state->payload->published_block_count.load(std::memory_order_seq_cst)) :
+                    -1,
+                staged_cores, drain_state_.drain_stage_done_mask.load(std::memory_order_relaxed),
+                drain_state_.drain_running_staged.load(std::memory_order_relaxed)
+            );
+        }
+    }
 #if SIMPLER_DFX
     // out param carries the PURE drain_stage_cores wall (build_payload + MMIO publish of
     // this thread's cores), isolating it from availability + stage_go handshake.
@@ -798,6 +981,37 @@ void SchedulerContext::handle_drain_mode(int32_t thread_idx, [[maybe_unused]] ui
         if (is_completed()) return;
         SPIN_WAIT_HINT();
     }
+    {
+        static std::atomic<int32_t> drain_aba_finalize_logs{0};
+        int32_t log_idx = drain_aba_finalize_logs.fetch_add(1, std::memory_order_relaxed);
+        if (log_idx < 64) {
+            int32_t staged_cores = 0;
+            if (slot_state->payload != nullptr) {
+                for (int w = 0; w < PTO2_EARLY_DISPATCH_CORE_MASK_WORDS; w++) {
+                    staged_cores += __builtin_popcountll(
+                        slot_state->payload->staged_core_mask[w].load(std::memory_order_seq_cst)
+                    );
+                }
+            }
+            LOG_INFO_V9(
+                "DRAIN_ABA_DIAG finalize sample=%d thread=%d task=%lld gated=%d stage_done=0x%x "
+                "running_staged=%d staged_cores=%d next_block=%d published=%d launch=%d drain_state=%d",
+                log_idx, thread_idx, static_cast<long long>(slot_state->task->task_id.raw), static_cast<int32_t>(gated),
+                drain_state_.drain_stage_done_mask.load(std::memory_order_relaxed),
+                drain_state_.drain_running_staged.load(std::memory_order_acquire), staged_cores,
+                slot_state->next_block_idx.load(std::memory_order_seq_cst),
+                slot_state->payload != nullptr ?
+                    static_cast<int32_t>(slot_state->payload->published_block_count.load(std::memory_order_seq_cst)) :
+                    -1,
+                slot_state->payload != nullptr ?
+                    static_cast<int32_t>(slot_state->payload->early_dispatch_launch_state.load(std::memory_order_seq_cst)) :
+                    -1,
+                slot_state->payload != nullptr ?
+                    static_cast<int32_t>(slot_state->payload->early_sync_drain_state.load(std::memory_order_seq_cst)) :
+                    -1
+            );
+        }
+    }
     if (gated) {
         // Seed the rendezvous with the running-slot cores staged across all threads; pending
         // cores advance it as they promote. maybe_rendezvous_ring (producer release) rings iff
@@ -824,7 +1038,26 @@ void SchedulerContext::handle_drain_mode(int32_t thread_idx, [[maybe_unused]] ui
     // ahead of drain completion and fail while running_slot_count is still incomplete. When
     // every block landed directly in a running slot, no pending promotion remains to retry it.
     if (gated) {
-        sched_->retry_sync_start_rendezvous_after_drain(*slot_state);
+        bool launched = sched_->retry_sync_start_rendezvous_after_drain(*slot_state);
+        static std::atomic<int32_t> drain_aba_rendezvous_logs{0};
+        int32_t log_idx = drain_aba_rendezvous_logs.fetch_add(1, std::memory_order_relaxed);
+        if (log_idx < 64) {
+            int32_t staged_cores = 0;
+            for (int w = 0; w < PTO2_EARLY_DISPATCH_CORE_MASK_WORDS; w++) {
+                staged_cores +=
+                    __builtin_popcountll(slot_state->payload->staged_core_mask[w].load(std::memory_order_seq_cst));
+            }
+            LOG_INFO_V9(
+                "DRAIN_ABA_DIAG rendezvous sample=%d thread=%d task=%lld launched=%d running_slot=%d "
+                "staged_cores=%d launch=%d early_dispatch=%d",
+                log_idx, thread_idx, static_cast<long long>(slot_state->task->task_id.raw),
+                static_cast<int32_t>(launched),
+                static_cast<int32_t>(slot_state->payload->running_slot_count.load(std::memory_order_seq_cst)),
+                staged_cores,
+                static_cast<int32_t>(slot_state->payload->early_dispatch_launch_state.load(std::memory_order_seq_cst)),
+                static_cast<int32_t>(slot_state->payload->early_dispatch_state.load(std::memory_order_seq_cst))
+            );
+        }
     } else {
         sched_->propagate_dispatch_fanin(*slot_state);
     }
diff --git a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_context.h b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_context.h
index a6f31617..31ce15c3 100644
--- a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_context.h
+++ b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_context.h
@@ -185,6 +185,7 @@ private:
     int32_t sched_thread_num_{0};
     int32_t aicpu_thread_num_{0};
     int32_t cores_total_num_{0};
+    bool drain_aba_test_mode_{false};
 
     // Cluster-ordered worker_id lists, populated by post_handshake_init().
     int32_t aic_worker_ids_[RUNTIME_MAX_WORKER]{};
diff --git a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_dispatch.cpp b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_dispatch.cpp
index c090f94c..67985a4c 100644
--- a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_dispatch.cpp
+++ b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_dispatch.cpp
@@ -421,8 +421,36 @@ void SchedulerContext::dispatch_shape(
                 }
                 int32_t available = is_mix ? selected_mix_clusters.count() : cores.count();
                 if (available < slot_state->logical_block_num) {
+                    static std::atomic<int32_t> drain_aba_enter_logs{0};
+                    int32_t log_idx = drain_aba_enter_logs.fetch_add(1, std::memory_order_relaxed);
+                    if (log_idx < 64) {
+                        LOG_INFO_V9(
+                            "DRAIN_ABA_DIAG enter_ready_drain sample=%d thread=%d task=%lld shape=%d local_available=%d "
+                            "block_num=%d",
+                            log_idx, thread_idx, static_cast<long long>(slot_state->task->task_id.raw),
+                            static_cast<int32_t>(shape), available, slot_state->logical_block_num
+                        );
+                    }
                     flush_publish();
-                    if (!enter_drain_mode(slot_state, slot_state->logical_block_num)) {
+                    bool drain_entered = enter_drain_mode(slot_state, slot_state->logical_block_num);
+                    static std::atomic<int32_t> drain_aba_enter_result_logs{0};
+                    int32_t result_log_idx = drain_aba_enter_result_logs.fetch_add(1, std::memory_order_relaxed);
+                    if (result_log_idx < 128) {
+                        PTO2TaskSlotState *pending = drain_state_.pending_task.load(std::memory_order_acquire);
+                        LOG_INFO_V9(
+                            "DRAIN_ABA_DIAG enter_ready_drain_result sample=%d thread=%d task=%lld entered=%d "
+                            "pending=%lld sync_start_pending=%d ack_mask=0x%x elected=%d",
+                            result_log_idx, thread_idx, static_cast<long long>(slot_state->task->task_id.raw),
+                            static_cast<int32_t>(drain_entered),
+                            pending != nullptr && pending->task != nullptr ?
+                                static_cast<long long>(pending->task->task_id.raw) :
+                                -1LL,
+                            drain_state_.sync_start_pending.load(std::memory_order_acquire),
+                            drain_state_.drain_ack_mask.load(std::memory_order_relaxed),
+                            drain_state_.drain_worker_elected.load(std::memory_order_relaxed)
+                        );
+                    }
+                    if (!drain_entered) {
                         disp_queues[static_cast<int32_t>(shape)].push(slot_state);
                     }
                     for (int rem = bi + 1; rem < got; rem++) {
@@ -1146,6 +1174,23 @@ int32_t SchedulerContext::resolve_and_dispatch(Runtime *runtime, int32_t thread_
 
         // Phase 2 drain check
         if (drain_state_.sync_start_pending.load(std::memory_order_acquire) != 0) {
+            if (deferred_release_count > 0) {
+                static std::atomic<int32_t> drain_aba_pre_drain_release_logs{0};
+                int32_t log_idx = drain_aba_pre_drain_release_logs.fetch_add(1, std::memory_order_relaxed);
+                if (log_idx < 64) {
+                    LOG_INFO_V9(
+                        "DRAIN_ABA_DIAG pre_drain_release sample=%d thread=%d count=%d",
+                        log_idx, thread_idx, deferred_release_count
+                    );
+                }
+                while (deferred_release_count > 0) {
+#if SIMPLER_SCHED_PROFILING
+                    (void)sched_->on_task_release(*deferred_release_slot_states[--deferred_release_count], thread_idx);
+#else
+                    sched_->on_task_release(*deferred_release_slot_states[--deferred_release_count]);
+#endif
+                }
+            }
 #if SIMPLER_DFX
             // The drain is otherwise a swimlane blind spot: the `continue` below skips
             // every phase record, and handle_drain_mode is uninstrumented. Time it here so
diff --git a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_types.h b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_types.h
index ff60af0f..28a222be 100644
--- a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_types.h
+++ b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_types.h
@@ -534,7 +534,10 @@ struct alignas(64) SyncStartDrainState {
     std::atomic<int32_t> drain_stage_go{0};          // 0=hold; 1=elected released parallel staging
     std::atomic<uint32_t> drain_stage_done_mask{0};  // bit per thread; all-set = all threads done staging
     std::atomic<int32_t> drain_running_staged{0};    // sum of running-slot cores staged (rendezvous seed)
-    int32_t _pad[7];
+    std::atomic<int32_t> drain_attempt_seq{0};        // test hook: generation for the current drain attempt
+    std::atomic<int32_t> drain_test_victim_armed{0};  // test hook: exactly one stale-attempt victim
+    std::atomic<int32_t> drain_test_stale_seen{0};    // test hook: latch first stale-attempt observation
+    int32_t _pad[4];
 };
 static_assert(sizeof(SyncStartDrainState) == 64);
 
diff --git a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/shared/runtime.cpp b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/shared/runtime.cpp
index a557cb30..a9cc7fe1 100644
--- a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/shared/runtime.cpp
+++ b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/shared/runtime.cpp
@@ -38,6 +38,7 @@ Runtime::Runtime() {
     dev.aicpu_allowed_cpu_count = 0;
     dev.aicpu_launch_count = 0;
     dev.serial_orch_sched = false;
+    dev.drain_aba_test_mode = false;
     dev.gm_sm_ptr_ = nullptr;
     dev.orch_args_storage_.clear();
     dev.prebuilt_arena_base_ = nullptr;

Complete ST code

The full ST code used for the reproducer is below. The paths are relative to the repository root.

Show full ST files

tests/st/a2a3/tensormap_and_ringbuffer/spmd_sync_start_drain_aba_hook/test_spmd_sync_start_drain_aba_hook.py

#!/usr/bin/env python3
# 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.
# -----------------------------------------------------------------------------------------------------------
"""Deterministic test-hook proof for the sync_start drain ABA window."""

import os

import torch
from simpler.task_interface import ArgDirection as D

from simpler_setup import SceneTestCase, TaskArgsBuilder, Tensor, scene_test

os.environ.setdefault("SIMPLER_DRAIN_ABA_TEST", "1")

FLOATS_PER_CACHE_LINE = 16
MIX_SLOTS = 3
SYNC_BLOCKS = 24
HOLDER_BLOCKS = 1


@scene_test(level=2, runtime="tensormap_and_ringbuffer")
class TestSpmdSyncStartDrainAbaHook(SceneTestCase):
    RTOL = 0
    ATOL = 0

    CALLABLE = {
        "orchestration": {
            "source": "kernels/orchestration/spmd_sync_start_drain_aba_hook_orch.cpp",
            "function_name": "aicpu_orchestration_entry",
            "signature": [D.INOUT, D.INOUT, D.IN],
        },
        "incores": [
            {
                "func_id": 0,
                "name": "SPMD_MIX_SLOW_AIC",
                "source": "../spmd_sync_start_drain_aba/kernels/aic/kernel_spmd_mix_pressure.cpp",
                "core_type": "aic",
                "signature": [D.INOUT, D.IN],
            },
            {
                "func_id": 1,
                "name": "SPMD_MIX_SLOW_AIV0",
                "source": "../spmd_sync_start_drain_aba/kernels/aiv/kernel_spmd_mix_pressure.cpp",
                "core_type": "aiv",
                "signature": [D.INOUT, D.IN],
            },
            {
                "func_id": 2,
                "name": "SPMD_MIX_SLOW_AIV1",
                "source": "../spmd_sync_start_drain_aba/kernels/aiv/kernel_spmd_mix_pressure.cpp",
                "core_type": "aiv",
                "signature": [D.INOUT, D.IN],
            },
        ],
    }

    CASES = [
        {
            "name": "Case1",
            "platforms": ["a2a3"],
            "config": {
                "aicpu_thread_num": 4,
                "block_dim": 24,
                "runtime_env": {"ring_heap": 32 * 1024 * 1024, "ring_task_window": 1024},
            },
            "params": {},
        }
    ]

    def generate_args(self, params):
        return TaskArgsBuilder(
            Tensor("output", torch.zeros(SYNC_BLOCKS * MIX_SLOTS * FLOATS_PER_CACHE_LINE, dtype=torch.float32)),
            Tensor("holder", torch.zeros(HOLDER_BLOCKS * MIX_SLOTS * FLOATS_PER_CACHE_LINE, dtype=torch.float32)),
            Tensor("scratch", torch.zeros(1024, dtype=torch.float32)),
        )

    def compute_golden(self, args, params):
        out = args.output
        for block_idx in range(SYNC_BLOCKS):
            for slot in range(MIX_SLOTS):
                out[(block_idx * MIX_SLOTS + slot) * FLOATS_PER_CACHE_LINE] = float(block_idx)


if __name__ == "__main__":
    SceneTestCase.run_module(__name__)

tests/st/a2a3/tensormap_and_ringbuffer/spmd_sync_start_drain_aba_hook/kernels/orchestration/spmd_sync_start_drain_aba_hook_orch.cpp

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

#include <stdint.h>

#include "pto_orchestration_api.h"

#define FUNC_SPMD_MIX_AIC 0
#define FUNC_SPMD_MIX_AIV0 1
#define FUNC_SPMD_MIX_AIV1 2

static constexpr int16_t HOLDER_BLOCKS = 1;
static constexpr int16_t SYNC_BLOCKS = 24;
static constexpr int64_t HOLDER_SPIN_ITERS = 50000000;

extern "C" {

__attribute__((visibility("default"))) PTO2OrchestrationConfig aicpu_orchestration_config(const L2TaskArgs &orch_args) {
    (void)orch_args;
    return PTO2OrchestrationConfig{.expected_arg_count = 3};
}

static MixedKernels mix_kernels() {
    MixedKernels mk;
    mk.aic_kernel_id = FUNC_SPMD_MIX_AIC;
    mk.aiv0_kernel_id = FUNC_SPMD_MIX_AIV0;
    mk.aiv1_kernel_id = FUNC_SPMD_MIX_AIV1;
    return mk;
}

static void submit_mix(
    const Tensor &out, const Tensor &scratch, int16_t block_num, int64_t base_cl, int64_t spin_iters, bool sync_start
) {
    L0TaskArgs args;
    args.add_inout(out);
    args.add_input(scratch);
    args.add_scalar(base_cl);
    args.add_scalar(spin_iters);
    args.launch_spec.set_block_num(block_num);
    args.launch_spec.set_require_sync_start(sync_start);
    rt_submit_task(mix_kernels(), args);
}

__attribute__((visibility("default"))) void aicpu_orchestration_entry(const L2TaskArgs &orch_args) {
    const Tensor &output = orch_args.tensor(0).ref();
    const Tensor &holder = orch_args.tensor(1).ref();
    const Tensor &scratch = orch_args.tensor(2).ref();

    submit_mix(holder, scratch, HOLDER_BLOCKS, 0, HOLDER_SPIN_ITERS, false);
    submit_mix(output, scratch, SYNC_BLOCKS, 0, 0, true);

    LOG_INFO_V9("[spmd_sync_start_drain_aba_hook] Submitted holder + sync_start drain reproducer");
}

}  // extern "C"

tests/st/a2a3/tensormap_and_ringbuffer/spmd_sync_start_drain_aba/kernels/aic/kernel_spmd_mix_pressure.cpp

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

/**
 * SPMD MIX pressure kernel (AIC): depends on a runtime-allocated scratch tensor,
 * spins to hold the cluster, then writes float(block_idx) to the output slot.
 *
 * Args: args[0] = output Tensor* (INOUT), args[1] = scratch Tensor* (INPUT),
 *       args[2] = base_cl, args[3] = spin_iters.
 */

#include <cstdint>
#include <pto/pto-inst.hpp>

#include "tensor.h"

#ifndef __gm__
#define __gm__
#endif

#ifndef __aicore__
#define __aicore__ [aicore]  // NOLINT(whitespace/braces)
#endif

#include "intrinsic.h"

static constexpr int32_t FLOATS_PER_CACHE_LINE = 16;
static constexpr int32_t SLOTS_PER_BLOCK = 3;

#ifdef PTO_CPUSTUB_HPP
#define dcci(...) \
    do {          \
    } while (0)
#endif
#ifndef SINGLE_CACHE_LINE
#define SINGLE_CACHE_LINE 0
#endif
#ifndef CACHELINE_OUT
#define CACHELINE_OUT 0
#endif

extern "C" __aicore__ void kernel_entry(__gm__ int64_t *args) {
    __gm__ Tensor *out_tensor = reinterpret_cast<__gm__ Tensor *>(args[0]);
    __gm__ Tensor *scratch_tensor = reinterpret_cast<__gm__ Tensor *>(args[1]);
    __gm__ float *out = reinterpret_cast<__gm__ float *>(out_tensor->buffer.addr) + out_tensor->start_offset;
    __gm__ float *scratch =
        reinterpret_cast<__gm__ float *>(scratch_tensor->buffer.addr) + scratch_tensor->start_offset;

    int32_t base_cl = static_cast<int32_t>(args[2]);
    int32_t spin_iters = static_cast<int32_t>(args[3]);
    int32_t block_idx = get_block_idx(args);
    int32_t offset = (base_cl + block_idx * SLOTS_PER_BLOCK + 0) * FLOATS_PER_CACHE_LINE;

    volatile float seed = scratch[block_idx & 1023];
    volatile int32_t acc = 0;
    for (int32_t i = 0; i < spin_iters; i++) {
        acc++;
    }

    out[offset] = static_cast<float>(block_idx) + seed * 0.0f + static_cast<float>(acc) * 0.0f;
    dcci(&out[offset], SINGLE_CACHE_LINE, CACHELINE_OUT);
}

tests/st/a2a3/tensormap_and_ringbuffer/spmd_sync_start_drain_aba/kernels/aiv/kernel_spmd_mix_pressure.cpp

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

/**
 * SPMD MIX pressure kernel (AIV): depends on a runtime-allocated scratch tensor,
 * spins to hold the cluster, then writes float(block_idx) to the AIV output slot.
 *
 * Args: args[0] = output Tensor* (INOUT), args[1] = scratch Tensor* (INPUT),
 *       args[2] = base_cl, args[3] = spin_iters.
 */

#include <cstdint>
#include <pto/pto-inst.hpp>

#include "tensor.h"

#ifndef __gm__
#define __gm__
#endif

#ifndef __aicore__
#define __aicore__ [aicore]  // NOLINT(whitespace/braces)
#endif

#include "intrinsic.h"

static constexpr int32_t FLOATS_PER_CACHE_LINE = 16;
static constexpr int32_t SLOTS_PER_BLOCK = 3;

#ifdef PTO_CPUSTUB_HPP
#define dcci(...) \
    do {          \
    } while (0)
#endif
#ifndef SINGLE_CACHE_LINE
#define SINGLE_CACHE_LINE 0
#endif
#ifndef CACHELINE_OUT
#define CACHELINE_OUT 0
#endif

extern "C" __aicore__ void kernel_entry(__gm__ int64_t *args) {
    __gm__ Tensor *out_tensor = reinterpret_cast<__gm__ Tensor *>(args[0]);
    __gm__ Tensor *scratch_tensor = reinterpret_cast<__gm__ Tensor *>(args[1]);
    __gm__ float *out = reinterpret_cast<__gm__ float *>(out_tensor->buffer.addr) + out_tensor->start_offset;
    __gm__ float *scratch =
        reinterpret_cast<__gm__ float *>(scratch_tensor->buffer.addr) + scratch_tensor->start_offset;

    int32_t base_cl = static_cast<int32_t>(args[2]);
    int32_t spin_iters = static_cast<int32_t>(args[3]);
    int32_t block_idx = get_block_idx(args);
    int32_t sub_block_id = get_sub_block_id(args);
    int32_t offset = (base_cl + block_idx * SLOTS_PER_BLOCK + 1 + sub_block_id) * FLOATS_PER_CACHE_LINE;

    volatile float seed = scratch[block_idx & 1023];
    volatile int32_t acc = 0;
    for (int32_t i = 0; i < spin_iters; i++) {
        acc++;
    }

    out[offset] = static_cast<float>(block_idx) + seed * 0.0f + static_cast<float>(acc) * 0.0f;
    dcci(&out[offset], SINGLE_CACHE_LINE, CACHELINE_OUT);
}

中文简要说明

这个问题本质上是 sync_start drain 协议里的可复用 barrier 没有 generation/attempt tag。owner 在资源不足时会清空 ack_mask 和 drain_worker_elected 让各线程 retry，但旧 attempt 中已经越过 ack barrier 的线程仍可能继续执行。这样它会把新 attempt 的空 election 状态当作旧 attempt 的状态来使用。

复现中，第一次 drain 三个线程都 ack，owner 看到 available=23 < block_num=24 后 reset。随后 thread1 作为旧 attempt 的 stale 线程继续原协议，发布 stage_go；而 thread0/thread2 已经在新 attempt 的 ack barrier 中等待 ack_mask=0x7。最终形成跨 generation 的 split wait：

thread1 等 stage_done == 0x7
thread0/thread2 等 ack_mask == 0x7
实际 ack_mask == 0x5

因此任务稳定超时并报 RuntimeError: run failed with code 507018。建议修复方向是在 drain ack/election/stage/finalize 全流程加入 generation 校验；资源不足 retry 时先推进 generation，再复用 ack/election 字段，使旧 attempt 线程能检测到 mismatch 并安全返回。

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    bugSomething isn't working

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions