[WIP] Add: SPMD paged attention example with dual-vector softmax - #6
Open
chenshengxin2026 wants to merge 4 commits into
Open
[WIP] Add: SPMD paged attention example with dual-vector softmax#6chenshengxin2026 wants to merge 4 commits into
chenshengxin2026 wants to merge 4 commits into
Conversation
Update all references in GitHub workflow skills, issue templates, and shared library docs to reflect the repo transfer. Co-authored-by: wcwxy <26245345+ChaoWao@users.noreply.github.com>
Two independent fixes to orchestration SO handling on AICPU: 1. Orch SO file creation split by platform. mkstemps (libdevice_orch_XXXXXX.so) ensures per-call uniqueness on sim where multiple workers may share a process, but is not always available on AICPU device libc. Added platform interface create_orch_so_file so sim uses mkstemps + fchmod(0755) and onboard uses pid-based naming + open(...,0755) — sufficient since only one runtime runs per device process. 2. Deferred dlclose/unlink from run() to deinit(). Closing the SO handle at the end of run() made it impossible to re-run the orchestrator through repeated calls into the same executor. The handle is kept until deinit, which then unlinks the file. Applied to a2a3 aicpu_build_graph, a2a3 tensormap_and_ringbuffer, and a5 tensormap_and_ringbuffer. Co-authored-by: wcwxy <26245345+ChaoWao@users.noreply.github.com>
Implements a complete paged attention kernel using SPMD parallelism under the tensormap_and_ringbuffer runtime. The pipeline consists of QK matmul, softmax prepare, PV matmul, and online update stages with dual AIV lanes processing 8-row sub-tiles each for the softmax and accumulation phases.
chenshengxin2026
force-pushed
the
add-spmd-paged-attention-example
branch
from
April 14, 2026 01:38
7ff6660 to
50b5e3b
Compare
chenshengxin2026
pushed a commit
that referenced
this pull request
Apr 15, 2026
…native-sys#563) Removes the fixed DIST_TASK_WINDOW_SIZE slot pool and the per-slot array DistWorker used to carry. At L3 the slot state lives entirely in the parent process's heap -- never crossed into child workers -- so the ring index L2 uses to address shmem descriptors buys us nothing here. Only the heap needs a pre-sized region for MAP_SHARED fork inheritance. - DistRing: - init() drops window_size; takes only (heap_bytes, timeout_ms). - alloc() returns a monotonic task id; no back-pressure on slot count, only on heap space. - Owns the slot state pool as std::deque<std::unique_ptr<SlotState>>. push_back never invalidates existing pointers, so slot_state(id) returns a pointer that stays valid for the slot's lifetime without holding the mutex past the lookup. - released_ and slot_heap_end_ become std::vector<>, grown via push_back on alloc, indexed directly by task id. - advance_last_alive_locked no longer needs to undo the released bit (entries aren't recycled within a run; reset_to_empty clears them all at drain). - New reset_to_empty(): drops all slot state and zeroes counters. DistOrchestrator::drain() calls it right after active_tasks_ hits 0 so each Worker.run() starts from task id 0 with bounded memory. - DistOrchestrator::init drops slots/num_slots params. slot_state(id) delegates to ring.slot_state(id) with a nullptr->throw guard. - DistScheduler::Config drops slots/num_slots; takes DistRing* and reads slot state via ring->slot_state(id) at every access site. - DistWorker drops the std::unique_ptr<SlotState[]> member; slot state is now entirely in allocator_. DistWorker::init() is a straight passthrough to allocator_/orchestrator_/scheduler_. - dist_types.h: remove DIST_TASK_WINDOW_SIZE constant. Tests: - test_dist_ring rewritten: drop window_size tests, add SlotAllocGrowsPastLegacyWindow (2048 allocs past the old 128 cap), SlotStateIsPointerStable (push_back doesn't invalidate refs), ResetToEmptyRequiresAllReleased, ResetToEmptyResetsCounters. - test_dist_orchestrator / test_dist_scheduler fixtures drop the std::unique_ptr<SlotState[]> member and access via a local S(id) helper that calls ring.slot_state(id). Docs: - orchestrator.md section 5 rewritten to describe the three resources DistRing now owns (task id, heap, slot state) and the end-of-run reset contract. - roadmap.md Dispatch internals bullet updated. Plan (local, gitignored): PR-I moved to "in review"; Allowed Exception #6 kept (explains why L3 doesn't need a shmem slot ring). No user-visible behaviour change: heap_ring_size still configurable via Worker ctor, OUTPUT auto-alloc / WaW tag semantics unchanged, back- pressure timeout still throws std::runtime_error on heap exhaustion.
chenshengxin2026
pushed a commit
that referenced
this pull request
Apr 21, 2026
…ative-sys#608) Introduce a one-shot cross-process mailbox class for parent-child bootstrap communication, independent of the task-mailbox protocol. Includes C++ implementation, nanobind Python bindings, and 7 UT cases covering in-process and fork-based cross-process scenarios. Design decisions: - Mailbox size: 4096 B (one page). HEADER_SIZE=64, ERROR_MSG_SIZE=1024, PTR_CAPACITY=376 — sufficient for all foreseeable chip buffer counts. - State machine: IDLE/SUCCESS/ERROR three states. Values 0/1/2 leave headroom for future intermediate states without serialization migration. - Memory ordering: aarch64 ldar/stlr inline asm (first, per codestyle #6), x86_64 compiler barrier, __atomic_load/store fallback — same pattern as WorkerThread mailbox in worker_manager.cpp. - Error message: strncpy with explicit null termination at size-1, compatible with L4 task-mailbox error message convention. Cross-process read hardening: - Ctor rejects max_buffer_count > CHIP_BOOTSTRAP_PTR_CAPACITY so the clamp invariant holds for every subsequent read. - buffer_ptrs() clamps the shared-memory count against max_buffer_count_ so a corrupted or premature read cannot overrun the pointer region. - error_message() uses strnlen(CHIP_BOOTSTRAP_ERROR_MSG_SIZE) instead of trusting the null-terminator in shared memory. Co-authored-by: wcwxy <26245345+ChaoWao@users.noreply.github.com>
chenshengxin2026
pushed a commit
that referenced
this pull request
Apr 21, 2026
…-native-sys#609) The C++ WorkerThread::dispatch_process already uses ldar/stlr (aarch64) and compiler-barriered plain store (x86_64) when it reads/writes the mailbox OFF_STATE word. The Python side of the handshake — the three worker loops (_sub_worker_loop, _chip_process_loop, _child_worker_loop), _chip_control, and the init/close paths — was using plain struct.pack_into("i", buf, _OFF_STATE, …) / unpack_from. On aarch64 that lets the state flip leak ahead of the preceding OFF_ERROR / OFF_ERROR_MSG writes, so a parent that observes TASK_DONE can read a stale error message. Design decisions (per .docs/l3-audit.md): 1. Exposure: add inline mailbox_load_i32 / mailbox_store_i32 in worker_bind.h and bind them as _mailbox_load_i32 / _mailbox_store_i32 on _task_interface. Underscore prefix keeps them out of task_interface.__all__ — only simpler.worker imports them. 2. ABI: aarch64 ldar/stlr first (per .claude/rules/codestyle.md #6), x86_64 second with __asm__ volatile("" ::: "memory") to stop the compiler from reordering across the TSO store, fallback to __atomic_{load,store} with ACQUIRE / RELEASE. 3. addr type: uint64_t. Python computes via ctypes.addressof(ctypes.c_char.from_buffer(buf)) + offset; C++ reinterpret_casts to volatile int32_t*. No void*. 4. Field ordering: all OFF_ERROR / OFF_ERROR_MSG / _CTRL_OFF_RESULT writes happen BEFORE the release-store of OFF_STATE, and all non-state reads happen AFTER the acquire-load of OFF_STATE. Every state access in python/simpler/worker.py now goes through the helper — 18 sites across the three loops, _chip_control, init, and close — so the invariant is mechanical. Adds tests/ut/py/test_worker/test_mailbox_atomics.py with four cases (roundtrip, cross-process visibility, payload-before-state ordering over 1000 fork iterations, refactored sub-worker dispatch). Existing L4 error-propagation tests still green.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
examples/a2a3/tensormap_and_ringbuffer/Test plan
python examples/scripts/run_example.py -k examples/a2a3/tensormap_and_ringbuffer/spmd_paged_attention/kernels -g examples/a2a3/tensormap_and_ringbuffer/spmd_paged_attention/golden.py -p sim