feat: 2O+2S Pipline

[TaoZQY]

Summary

This PR adds the strategy2 2O + 2S pipeline path for the a2a3 tensormap runtime.

The change is split into two parts:

• Add pipeline_strategy plumbing through CallConfig, Python bindings, remote wire encoding, simpler_run, and Runtime launch descriptors.
• Add the a2a3 strategy2 runtime implementation that splits the O-side submit path into two orchestrator stages.

O-side Pipeline Split

Strategy2 uses 4 AICPU threads:

• S0/S1: scheduler threads
• O0: primary orchestrator thread
• O1: submit-commit pipeline worker

The original O-side submit flow mixed several responsibilities in one orchestrator thread:

1. allocate task/payload/output metadata
2. compute fanin dependencies
3. register TensorMap outputs
4. publish task to scheduler wiring or ready queue

This PR splits that flow:

O0: prepare and enqueue

O0 keeps the orchestration/front-half work:

• run normal orchestration function
• allocate task slot and output layout
• initialize PTO2TaskPayload
• preserve explicit deps, tensor arg types, kernel IDs, manual-scope state
• enqueue a PTO2SubmitCommitRecord into the submit pipeline queue

O1: dependency commit and publish

O1 runs run_submit_pipeline_worker() and handles the back-half:

• compute deferred fanin dependencies from the prepared payload
• register TensorMap outputs
• handle scope-end records
• publish tasks to scheduler

For no-fanin tasks, O1 directly publishes to the ready queue via publish_ready_no_fanin(), avoiding the wiring queue path.

For tasks with fanin, O1 pushes into the scheduler wiring queue and signals scheduler drain when needed.

Behavior

• Baseline behavior is unchanged when pipeline_strategy is unset.
• Strategy2 is enabled by pipeline_strategy=2 or SIMPLER_PIPELINE_STRATEGY=2.
• If the requested AICPU thread count cannot support 2S + 2O, runtime falls back to baseline layout.
• SIMPLER_PIPELINE_DEFER_SUBMIT=0/false/off/no can disable deferred submit for debugging.

Results

Summary

Case	2O+2S Total Δ%	2O+2S Orch Δ%
alternating_matmul_add Case1	-20.76%	-20.73%
benchmark_bgemm Case0	-24.98%	-27.41%
paged_attention_unroll Case1	-7.35%	-14.13%
paged_attention_unroll Case2	-7.90%	-0.52%
paged_attention_unroll_manual_scope Case1	-6.67%	-2.01%
paged_attention_unroll_manual_scope Case2	-6.89%	+8.64%
batch_paged_attention Case1	-25.70%	-32.43%
spmd_paged_attention Case1	-0.51%	0.00%
spmd_paged_attention Case2	-1.37%	+0.93%
qwen14b-StressBatch16Seq3500	-6.24%	-23.81%

Note: A negative value for the change in the table indicates a reduction in time consumption relative to the baseline, which means improved performance.

Qwen14b-decode

Case:StressBatch16Seq3500

Metric	Base Avg	S2 Avg	S2 Avg Δ	S2 Avg Δ%	S2 Trimmed Avg	Trimmed Δ	Trimmed Δ%
Total / Effective	2237.2	2097.6	-139.6	-6.24%	2095.8	-141.4	-6.32%
Orch	530.0	403.8	-126.2	-23.81%	402.7	-127.3	-24.02%
Sched	2235.5	2096.4	-139.1	-6.22%	2094.6	-140.9	-6.30%

9 case Benchmark

Case	Base Total	S2 Total	Total Δ	Total Δ%	Base Orch	S2 Orch	Orch Δ	Orch Δ%
alternating_matmul_add Case1	882.3	699.1	-183.2	-20.76%	874.5	693.2	-181.3	-20.73%
benchmark_bgemm Case0	845.4	634.2	-211.2	-24.98%	764.6	555.0	-209.6	-27.41%
paged_attention_unroll Case1	1246.2	1154.6	-91.6	-7.35%	909.5	781.0	-128.5	-14.13%
paged_attention_unroll Case2	649.4	598.1	-51.3	-7.90%	420.9	418.7	-2.2	-0.52%
paged_attention_unroll_manual_scope Case1	1233.5	1151.2	-82.3	-6.67%	766.3	750.9	-15.4	-2.01%
paged_attention_unroll_manual_scope Case2	636.1	592.3	-43.8	-6.89%	372.6	404.8	+32.2	+8.64%
batch_paged_attention Case1	4150.3	3083.8	-1066.5	-25.70%	3251.1	2196.8	-1054.3	-32.43%
spmd_paged_attention Case1	1309.1	1302.4	-6.7	-0.51%	11.8	11.8	0.0	0.00%
spmd_paged_attention Case2	699.1	689.5	-9.6	-1.37%	10.7	10.8	+0.1	+0.93%

Test

SIMPLER_PIPELINE_STRATEGY=2 
SIMPLER_AICPU_THREAD_NUM=4 
./tools/benchmark_rounds.sh -p a2a3 -r tensormap_and_ringbuffer -d 0 -n 100

[coderabbitai]

Important

Review skipped

Draft detected.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro

Run ID: e0e63b83-1ae0-42b1-b13f-8d6c58f49631

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[gemini-code-assist]

Code Review

This pull request introduces a new multi-threaded pipeline strategy (S2_O2_SPLIT_CTRL_STRATEGY2) that defers dependency resolution, TensorMap registration, and task publication to a secondary orchestrator worker thread. The review feedback identifies several critical concurrency and synchronization issues in this new pipeline implementation. These include a potential data race on TensorMap during synchronous fallback, a potential deadlock caused by not updating the producer_blocked state when the pipeline worker is blocked, and multiple potential permanent hangs in spin loops (submit_pipeline_queue_push and flush_submit_pipeline) if a fatal error occurs.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

When the task has too many explicit dependencies or metadata serialization fails, the orchestrator falls through to the synchronous submit path on the primary thread (O0). This allows O0 to access the shared tensor_map concurrently with the pipeline worker thread (O1) processing deferred tasks, leading to data races. Since TensorMap is accessed by multiple threads, its operations must be protected by a mutex to prevent data races, even if current usage patterns attempt to avoid overlap.

    std::lock_guard<std::mutex> lock(orch->tensor_map_mutex);
    PTO2FaninBuilder fanin_builder(orch, orch->rings[ring_id].fanin_pool, next_fanin_seen_epoch(orch));

References

1. If a component (e.g., TensorMap) is accessed by multiple threads, its operations should be protected by a mutex to prevent data races, even if current usage patterns do not show overlap.

[gemini-code-assist]

When the pipeline worker thread (O1) is blocked pushing to the scheduler's wiring queue, it does not set producer_blocked to 1. This prevents the scheduler's deadlock detector from realizing that the producer is blocked, which can lead to permanent deadlocks. Please set and reset producer_blocked around the blocking push loop, matching the baseline's pattern.

    } else {
        if (!record.scheduler->wiring.queue.push(record.slot_state)) {
            record.scheduler->wiring.producer_blocked.store(1, std::memory_order_release);
            while (!record.scheduler->wiring.queue.push(record.slot_state)) {
                publish_spins++;
                if (orch->sm_header->orch_error_code.load(std::memory_order_acquire) != PTO2_ERROR_NONE) {
                    orch->fatal = true;
                    record.scheduler->wiring.producer_blocked.store(0, std::memory_order_release);
                    return;
                }
                SPIN_WAIT_HINT();
            }
            record.scheduler->wiring.producer_blocked.store(0, std::memory_order_release);
        }
        signal_scheduler_drain_if_needed(orch, record);
    }

[gemini-code-assist]

If the submit pipeline queue is full, the primary orchestrator thread (O0) spins indefinitely in submit_pipeline_queue_push. If the pipeline worker thread (O1) or the scheduler encounters a fatal error and aborts, O1 will stop popping, causing O0 to hang forever. Please pass orch to submit_pipeline_queue_push and check for fatal errors inside the spin loop to prevent permanent hangs.

static void submit_pipeline_queue_push(PTO2OrchestratorState *orch, PTO2SubmitPipelineQueue &queue, const PTO2SubmitCommitRecord &record) {
    uint64_t seq = queue.tail.load(std::memory_order_acquire);
    int32_t slot = static_cast<int32_t>(seq % PTO2_SUBMIT_PIPELINE_QUEUE_CAP);
    while (queue.slot_state[slot].load(std::memory_order_acquire) != 0) {
        if (orch->fatal || orch->sm_header->orch_error_code.load(std::memory_order_acquire) != PTO2_ERROR_NONE) {
            orch->fatal = true;
            return;
        }
        SPIN_WAIT_HINT();
    }
    queue.records[slot] = record;
    queue.slot_state[slot].store(1, std::memory_order_release);
    queue.tail.store(seq + 1, std::memory_order_release);
}

[gemini-code-assist]

Update the call to submit_pipeline_queue_push to pass orch as the first argument, matching the updated signature that checks for fatal errors.

    submit_pipeline_queue_push(orch, orch->submit_pipeline_queues[0], record);

[gemini-code-assist]

Update the call to submit_pipeline_queue_push to pass orch as the first argument, matching the updated signature that checks for fatal errors.

    submit_pipeline_queue_push(orch, orch->submit_pipeline_queues[0], record);

[gemini-code-assist]

If the pipeline worker thread (O1) aborts due to a fatal error, submit_pipeline_completed will never reach target, causing the primary orchestrator thread (O0) to hang indefinitely in flush_submit_pipeline. Please check for fatal errors inside the spin loop to break out safely.

    while (submit_pipeline_completed.load(std::memory_order_acquire) < target) {
        if (fatal || sm_header->orch_error_code.load(std::memory_order_acquire) != PTO2_ERROR_NONE) {
            fatal = true;
            break;
        }
        SPIN_WAIT_HINT();
    }