feat: support pipelined orchestration and scheduling across tokens and requests#1379
feat: support pipelined orchestration and scheduling across tokens and requests#1379TaoZQY wants to merge 17 commits into
Conversation
|
Important Review skippedAuto incremental reviews are disabled on this repository. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Plus Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
📝 WalkthroughWalkthroughThis PR adds concurrent streaming request sessions for hierarchical workers, including bounded admission, per-request token delivery, cancellation, prepared host/device execution, epoch-based HostGraph publication, arena-bank coordination, concurrency synchronization, paged-attention coverage, and Perfetto profiling tools. ChangesStreaming HostGraph execution
Estimated code review effort: 5 (Critical) | ~120 minutes Possibly related PRs
Sequence Diagram(s)sequenceDiagram
participant RequestSession
participant Worker
participant HostRequestAdmissionService
participant ChipWorker
participant AicpuExecutor
RequestSession->>Worker: submit request
Worker->>HostRequestAdmissionService: prepare host request
HostRequestAdmissionService->>ChipWorker: prepare request
Worker->>ChipWorker: execute prepared request
ChipWorker->>AicpuExecutor: launch device execution
AicpuExecutor-->>RequestSession: publish streamed token
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
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. Comment |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces Streaming Request Sessions to enable non-blocking request admission and Host O / Device S pipelining on hierarchical Workers. It implements RequestSession and RequestStream in Python, adds support for concurrent runs and async HostGraph building in the C++ runtime and scheduler, and introduces split request preparation on the ChipWorker. Additionally, comprehensive integration and unit tests are added. The review feedback suggests replacing a hardcoded magic number in the shared memory size configuration with explicit constants to improve maintainability and prevent layout synchronization issues.
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.
| _CONTROL_OFF_RESPONSE = _CONTROL_OFF_REQUEST + _REQUEST.size | ||
| CONTROL_BLOCK_SIZE = _CONTROL_OFF_RESPONSE + _RESPONSE.size | ||
| CONTROL_SHM_SIZE = max(4096, CONTROL_BLOCK_SIZE) | ||
| CONTROL_SHM_SIZE = max(65536, CONTROL_BLOCK_SIZE) |
There was a problem hiding this comment.
The hardcoded value 65536 makes the code harder to maintain as it obscures the origin of this number. It appears to be derived from HOST_REQUEST_CONTROL_OFFSET and HOST_REQUEST_CONTROL_BYTES in the new request_session.py file.
To improve clarity and avoid this magic number, please define these values as constants here. This makes the relationship between the shared memory size and the control region it contains explicit and ensures they stay in sync if the layout changes.
| CONTROL_SHM_SIZE = max(65536, CONTROL_BLOCK_SIZE) | |
| HOST_REQUEST_CONTROL_OFFSET = 4096 | |
| HOST_REQUEST_CONTROL_BYTES = 60 * 1024 | |
| CONTROL_SHM_SIZE = max(HOST_REQUEST_CONTROL_OFFSET + HOST_REQUEST_CONTROL_BYTES, CONTROL_BLOCK_SIZE) |
139fe3a to
7e31f63
Compare
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 18
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@docs/request-session.md`:
- Around line 69-79: Update the example’s request preparation flow after
emitter.prepare_host_request to call append_host_graph_prepared_request_args
with chip_args and request_id before orch.submit_next_level, ensuring the
submission is recognized as a prepared host-graph request and uses the prepared
arena bank.
In `@python/simpler/l3_l2_orch_comm.py`:
- Around line 214-220: Ensure timeout handling cannot leave either shared-memory
control lane permanently in DONE: in python/simpler/l3_l2_orch_comm.py lines
214-220, update the L3-L2 orchestration client flow around _read_response and
the polling loop to invalidate and recreate the client after a timeout, or
implement an equivalent race-safe acknowledged reset; apply the same recovery
contract to HostRequestAdmissionClient in python/simpler/request_session.py
lines 114-122. Preserve normal successful completion and IDLE transition
behavior.
In `@python/simpler/request_session.py`:
- Around line 395-408: Update close() to reject calls from a dispatcher thread
by checking _owns_current_thread() before enqueueing _STOP or releasing
ownership, and fail fast when true. Preserve normal shutdown and ownership
release for callers outside the dispatcher threads.
- Around line 246-267: The _emit method currently completes the stream
immediately when final=True, allowing later failures in Worker.run processing to
be reported as success. Remove or defer the _finish(None) call from _emit, and
perform successful termination only at the existing completion point after
Worker.run returns in _dispatch_loop around line 454; preserve cancellation and
terminal-error propagation.
- Around line 374-382: Update the request submission sequence around
_RequestWork and _live_work so each work item is inserted into _live_work before
calling self._requests.put_nowait(work). Preserve the existing
RequestBackpressureError behavior, and ensure a full queue does not leave an
unsubmitted work item registered in _live_work.
In `@python/simpler/worker.py`:
- Around line 4067-4072: Update both shared-memory cleanup sites in
python/simpler/worker.py at lines 4067-4072 and 3861-3866: make close and unlink
independent operations so a SharedMemory.close() failure does not prevent
unlink(). Preserve the bootstrap rollback’s original exception by suppressing
cleanup errors there, and retain best-effort cleanup during normal teardown.
- Around line 3835-3860: Serialize the per-worker initialization in
_ensure_l3_l2_orch_comm by guarding the readiness check, control SharedMemory
creation, service initialization, and client publication with a lock. Re-check
_l3_l2_orch_comm_ready after acquiring the lock, return the existing client when
another dispatcher completed initialization, and ensure only one child service
and dictionary entry are created per worker.
- Around line 459-476: Update _prepared_request_id_from_blob to unpack the
fourth trailer scalar following token_magic, then reject the trailer unless the
prepared request ID and stream request ID match. Preserve the existing
magic-value validation and return the validated request ID only when both IDs
agree.
In `@src/a2a3/runtime/host_build_graph/host/runtime_maker.cpp`:
- Around line 1142-1164: Update wait_for_device_epoch to poll the device’s
orch_error_code and sched_error_code alongside the epoch. Return immediately
with an error when either field reports a fatal error, while retaining the
existing 30-second deadline only as the deadlock fallback.
In `@src/common/platform/onboard/host/c_api_shared.cpp`:
- Around line 680-682: Update both exceptional cleanup sites in
src/common/platform/onboard/host/c_api_shared.cpp: at lines 680-682, mirror the
prepare-failure cleanup by invoking validate_runtime_impl() before destroying r;
at lines 721-722, invoke validate_runtime_impl() before the runtime scope guard
destroys r. Ensure both catch paths reconcile prepared Runtime resources before
returning or scope cleanup.
- Around line 73-85: Update capture_thread_context(), bind_thread_context(), and
unbind_thread_context() to preserve the current runner together with the
selected arena-bank TLS value. Store both pieces of context in the captured
state, restore the runner and bank when binding the async thread, and clear or
restore both values when unbinding, while retaining the existing attach failure
cleanup behavior.
In `@src/common/worker/chip_worker.cpp`:
- Around line 391-417: Update the preparation flow around prepared_requests_ to
reserve request_id under prepared_mu_ before waiting for or invoking
prepare_request_fn_. Track the reservation as in-progress, then atomically
publish the completed prepared state while preserving duplicate-ID rejection.
Ensure every preparation, allocation, and publication failure releases the arena
bank and removes the reservation; add a regression test covering concurrent
preparations with the same request ID.
In
`@tests/st/a2a3/host_build_graph/batch_paged_attention_three_layer/kernels/aic/aic_qk_matmul.cpp`:
- Around line 125-144: Reject invalid dimensions before dispatch and defensively
inside each kernel entrypoint: in
tests/st/a2a3/host_build_graph/batch_paged_attention_three_layer/kernels/aic/aic_qk_matmul.cpp:125-144,
validate positive batch_count, divisibility, and exact supported
query-tile/block-size/head-dimension tuples before qk_matmul_batch_impl; apply
the equivalent combined validation in aic_pv_matmul.cpp:121-137 before its
fixed-template dispatch. In aiv_online_update.cpp:211-230, accept only (16,16),
(16,128), and (64,128), and in aiv_softmax_prepare.cpp:184-200 accept only
(16,16), (16,128), and (64,64); reject unsupported inputs at launch and repeat
the checks in each kernel to prevent invalid dispatch and out-of-bounds
execution.
In
`@tests/st/a2a3/host_build_graph/batch_paged_attention_three_layer/kernels/orchestration/paged_attention_three_layer_orch.cpp`:
- Around line 83-90: Move the layer barrier capacity validation in the
orchestration routine so it runs before any tasks are submitted for a layer.
Reject the layer when q_loop * num_chunks exceeds MAX_LAYER_DONE_DEPS, returning
through the existing failure path before creating partial dependencies. Preserve
normal submission and barrier/final-epoch behavior for valid layer sizes.
In
`@tests/st/a2a3/host_build_graph/batch_paged_attention_three_layer/test_batch_paged_attention_three_layer_l3.py`:
- Around line 223-242: Update the request synchronization around _relay_tokens
and request_a_submitted so the event is set only when A’s first token (sequence
1) is received and published, rather than immediately after submit. Keep B’s
session.submit after waiting for that event, preserving the existing stagger
only after A’s first publication.
In `@tests/ut/py/test_worker/test_request_session.py`:
- Around line 69-85: Update test_emit_waits_until_consumer_receives_token to add
a producer-entry synchronization event set immediately before
emitter.emit("token-1"), then wait for that event before asserting producer_done
is unset. Preserve the existing token consumption, completion, stream cleanup,
and thread-join assertions.
In `@tools/cross_request_host_device_perfetto.py`:
- Around line 85-131: Update parse_log and map_host_invocations so
request-to-invocation mapping does not depend on token arrival order. Parse
HOST_REQUEST_PREPARED records and associate each prepared request ID with its
corresponding simpler_prepare_request invocation, then assign the remaining
request to the sole simpler_run invocation. Preserve validation that exactly one
normal invocation and one prepare invocation exist for each later request, and
reject missing or ambiguous mappings.
- Around line 187-191: Update the delivery validation around received,
delivered, and one_to_one to count every (request, seq) occurrence with Counter
before declaring strict one-to-one delivery. Require the Counter values to
match, while retaining the existing timestamp-map calculation for shared-key
latency and nonnegative delays.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 2f2b214f-ad20-412f-b4b5-74851e9b0009
📒 Files selected for processing (53)
README.mddocs/dynamic-linking.mddocs/hierarchical_level_runtime.mddocs/request-session.mdpython/bindings/task_interface.cpppython/simpler/l3_l2_orch_comm.pypython/simpler/request_session.pypython/simpler/worker.pysrc/a2a3/platform/onboard/host/device_runner.cppsrc/a2a3/platform/sim/host/device_runner.cppsrc/a2a3/runtime/host_build_graph/aicpu/aicpu_executor.cppsrc/a2a3/runtime/host_build_graph/host/runtime_maker.cppsrc/a2a3/runtime/host_build_graph/orchestration/common.cppsrc/a2a3/runtime/host_build_graph/orchestration/pto_orchestration_api.hsrc/a2a3/runtime/host_build_graph/runtime/host_graph_epoch.hsrc/a2a3/runtime/host_build_graph/runtime/host_graph_token_stream.hsrc/a2a3/runtime/host_build_graph/runtime/orchestrator_core/pto_orchestrator.cppsrc/a2a3/runtime/host_build_graph/runtime/orchestrator_core/pto_runtime2.cppsrc/a2a3/runtime/host_build_graph/runtime/pto_runtime2.hsrc/a2a3/runtime/host_build_graph/runtime/pto_shared_memory.hsrc/a2a3/runtime/host_build_graph/runtime/runtime.hsrc/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_cold_path.cppsrc/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_context.hsrc/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_dispatch.cppsrc/a2a3/runtime/host_build_graph/runtime/shared/pto_shared_memory.cppsrc/a2a3/runtime/host_build_graph/runtime/shared/runtime.cppsrc/common/hierarchical/orchestrator.cppsrc/common/hierarchical/orchestrator.hsrc/common/hierarchical/scope.cppsrc/common/hierarchical/scope.hsrc/common/hierarchical/tensormap.cppsrc/common/hierarchical/tensormap.hsrc/common/log/include/common/strace.hsrc/common/platform/include/common/host_api.hsrc/common/platform/onboard/host/c_api_shared.cppsrc/common/platform/onboard/host/device_runner_base.cppsrc/common/platform/onboard/host/device_runner_base.hsrc/common/platform/sim/host/c_api_shared.cppsrc/common/worker/chip_worker.cppsrc/common/worker/chip_worker.hsrc/common/worker/pto_runtime_c_api.htests/st/a2a3/host_build_graph/batch_paged_attention_three_layer/kernels/aic/aic_pv_matmul.cpptests/st/a2a3/host_build_graph/batch_paged_attention_three_layer/kernels/aic/aic_qk_matmul.cpptests/st/a2a3/host_build_graph/batch_paged_attention_three_layer/kernels/aiv/aiv_online_update.cpptests/st/a2a3/host_build_graph/batch_paged_attention_three_layer/kernels/aiv/aiv_softmax_prepare.cpptests/st/a2a3/host_build_graph/batch_paged_attention_three_layer/kernels/orchestration/paged_attention_three_layer_orch.cpptests/st/a2a3/host_build_graph/batch_paged_attention_three_layer/test_batch_paged_attention_three_layer.pytests/st/a2a3/host_build_graph/batch_paged_attention_three_layer/test_batch_paged_attention_three_layer_l3.pytests/ut/cpp/common/test_trb_runtime_temp_buffer.cpptests/ut/cpp/hierarchical/test_scope.cpptests/ut/py/test_worker/test_request_session.pytools/cross_request_host_device_perfetto.pytools/host_orch_perfetto.py
| ```python | ||
| emitter.prepare_host_request( | ||
| worker_id=0, | ||
| request_id=request_id, | ||
| callable_handle=callable_handle, | ||
| args=chip_args, | ||
| config=config, | ||
| arena_bank=1, | ||
| ) | ||
| orch.submit_next_level(callable_handle, chip_args, config, worker=0) | ||
| ``` |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Mark chip_args as a prepared request before submission.
The example calls prepare_host_request but never appends HOST_GRAPH_PREPARED_REQUEST_MAGIC. Consequently, the chip loop uses normal run_from_blob, leaving the prepared request and arena bank stranded. Add append_host_graph_prepared_request_args(chip_args, request_id) before submit_next_level.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@docs/request-session.md` around lines 69 - 79, Update the example’s request
preparation flow after emitter.prepare_host_request to call
append_host_graph_prepared_request_args with chip_args and request_id before
orch.submit_next_level, ensuring the submission is recognized as a prepared
host-graph request and uses the prepared arena bank.
| while _mailbox_load_i32(self._state_addr) != _STATE_DONE: | ||
| if time.monotonic() >= deadline: | ||
| raise TimeoutError("L3-L2 orch comm client timed out waiting for DONE") | ||
| time.sleep(_POLL_INTERVAL_S) | ||
|
|
||
| response = self._read_response() | ||
| _mailbox_store_i32(self._state_addr, _STATE_IDLE) |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | 🏗️ Heavy lift
Timeouts permanently poison both shared-memory control lanes. Once READY is published, either service may complete after the caller times out, leaving DONE with no client transition back to IDLE.
python/simpler/l3_l2_orch_comm.py#L214-L220: invalidate and recreate the orchestration client after timeout, or implement a race-safe acknowledged reset.python/simpler/request_session.py#L114-L122: apply the same recovery contract toHostRequestAdmissionClient.
📍 Affects 2 files
python/simpler/l3_l2_orch_comm.py#L214-L220(this comment)python/simpler/request_session.py#L114-L122
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@python/simpler/l3_l2_orch_comm.py` around lines 214 - 220, Ensure timeout
handling cannot leave either shared-memory control lane permanently in DONE: in
python/simpler/l3_l2_orch_comm.py lines 214-220, update the L3-L2 orchestration
client flow around _read_response and the polling loop to invalidate and
recreate the client after a timeout, or implement an equivalent race-safe
acknowledged reset; apply the same recovery contract to
HostRequestAdmissionClient in python/simpler/request_session.py lines 114-122.
Preserve normal successful completion and IDLE transition behavior.
| def _emit(self, item: Any, *, final: bool = False) -> None: | ||
| consumed = threading.Event() | ||
| with self._finish_lock: | ||
| if self._done.is_set(): | ||
| if isinstance(self._terminal_error, RequestCancelledError): | ||
| return | ||
| raise RuntimeError(f"request {self.request_id} stream is already complete") | ||
| if self._pending_consumed is not None: | ||
| raise RuntimeError(f"request {self.request_id} already has an outstanding item") | ||
| self._pending_consumed = consumed | ||
| self._items.put(_StreamItem(item, consumed)) | ||
| consumed.wait() | ||
| with self._finish_lock: | ||
| if self._pending_consumed is consumed: | ||
| self._pending_consumed = None | ||
| if self._done.is_set(): | ||
| if isinstance(self._terminal_error, RequestCancelledError): | ||
| return | ||
| if self._terminal_error is not None: | ||
| raise self._terminal_error | ||
| if final: | ||
| self._finish(None) |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
Do not report terminal success before Worker.run finishes.
final=True calls _finish(None) inside the orchestration callback. If the subsequent drain, validation, or copy-back fails, _dispatch_loop cannot replace that success because _finish is already complete; users observe a successful stream. Defer successful termination to Line 454 after Worker.run returns.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@python/simpler/request_session.py` around lines 246 - 267, The _emit method
currently completes the stream immediately when final=True, allowing later
failures in Worker.run processing to be reported as success. Remove or defer the
_finish(None) call from _emit, and perform successful termination only at the
existing completion point after Worker.run returns in _dispatch_loop around line
454; preserve cancellation and terminal-error propagation.
| stream = RequestStream(request_id, self) | ||
| work = _RequestWork(request, request_id, config, stream, threading.Event()) | ||
| try: | ||
| self._requests.put_nowait(work) | ||
| except queue.Full as exc: | ||
| raise RequestBackpressureError( | ||
| f"RequestSession pending queue is full (request_id={request_id})" | ||
| ) from exc | ||
| self._live_work[request_id] = work |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Register live work before publishing it to the dispatcher queue.
A dispatcher can consume and retire work immediately after put_nowait, before Line 382 inserts it. The completed request is then inserted permanently into _live_work, preventing reuse of its ID.
Proposed ordering fix
stream = RequestStream(request_id, self)
work = _RequestWork(request, request_id, config, stream, threading.Event())
+ self._live_work[request_id] = work
try:
self._requests.put_nowait(work)
except queue.Full as exc:
+ self._live_work.pop(request_id, None)
raise RequestBackpressureError(
f"RequestSession pending queue is full (request_id={request_id})"
) from exc
- self._live_work[request_id] = work📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| stream = RequestStream(request_id, self) | |
| work = _RequestWork(request, request_id, config, stream, threading.Event()) | |
| try: | |
| self._requests.put_nowait(work) | |
| except queue.Full as exc: | |
| raise RequestBackpressureError( | |
| f"RequestSession pending queue is full (request_id={request_id})" | |
| ) from exc | |
| self._live_work[request_id] = work | |
| stream = RequestStream(request_id, self) | |
| work = _RequestWork(request, request_id, config, stream, threading.Event()) | |
| self._live_work[request_id] = work | |
| try: | |
| self._requests.put_nowait(work) | |
| except queue.Full as exc: | |
| self._live_work.pop(request_id, None) | |
| raise RequestBackpressureError( | |
| f"RequestSession pending queue is full (request_id={request_id})" | |
| ) from exc |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@python/simpler/request_session.py` around lines 374 - 382, Update the request
submission sequence around _RequestWork and _live_work so each work item is
inserted into _live_work before calling self._requests.put_nowait(work).
Preserve the existing RequestBackpressureError behavior, and ensure a full queue
does not leave an unsubmitted work item registered in _live_work.
| def close(self) -> None: | ||
| should_stop = False | ||
| with self._state_lock: | ||
| if self._accepting: | ||
| self._accepting = False | ||
| should_stop = True | ||
| if should_stop: | ||
| for _thread in self._threads: | ||
| self._requests.put(_STOP) | ||
| current_ident = threading.get_ident() | ||
| for thread in self._threads: | ||
| if thread.ident != current_ident: | ||
| thread.join() | ||
| self._worker._release_request_session(self) |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Reject close() from a dispatcher thread.
With a full queue, the sole dispatcher can block forever while inserting _STOP. Even without a full queue, ownership is released before that dispatcher exits, allowing another session or direct run to overlap it. Fail fast when _owns_current_thread() is true, or defer ownership release until the loop exits.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@python/simpler/request_session.py` around lines 395 - 408, Update close() to
reject calls from a dispatcher thread by checking _owns_current_thread() before
enqueueing _STOP or releasing ownership, and fail fast when true. Preserve
normal shutdown and ownership release for callers outside the dispatcher
threads.
| constexpr uint64_t IN_CORE_BATCH = 16; | ||
| constexpr uint32_t MAX_LAYER_DONE_DEPS = 64; | ||
| uint64_t num_chunks = (batch + IN_CORE_BATCH - 1) / IN_CORE_BATCH; | ||
| PTO2TaskId prev_layer_done = PTO2TaskId::invalid(); | ||
|
|
||
| for (uint64_t layer = 0; layer < layer_count; layer++) { | ||
| PTO2TaskId layer_done_deps[MAX_LAYER_DONE_DEPS]; | ||
| uint32_t layer_done_dep_count = 0; |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Reject oversized layer barriers before submitting any tasks.
When q_loop * num_chunks > 64, the current check runs only after the 65th chunk has been submitted. Returning then leaves a partial layer without its barrier or final epoch, which can strand the streaming request.
Proposed fix
constexpr uint64_t IN_CORE_BATCH = 16;
constexpr uint32_t MAX_LAYER_DONE_DEPS = 64;
uint64_t num_chunks = (batch + IN_CORE_BATCH - 1) / IN_CORE_BATCH;
+ uint64_t deps_per_layer = q_loop * num_chunks;
+ if (deps_per_layer > MAX_LAYER_DONE_DEPS) {
+ LOG_ERROR(
+ "batch_paged_attention_three_layer: layer requires %" PRIu64
+ " barrier dependencies; maximum is %u",
+ deps_per_layer, MAX_LAYER_DONE_DEPS
+ );
+ return;
+ }
PTO2TaskId prev_layer_done = PTO2TaskId::invalid();
@@
- if (!chunk_done.is_valid() || layer_done_dep_count >= MAX_LAYER_DONE_DEPS) {
- LOG_ERROR("batch_paged_attention_three_layer: layer barrier capacity exceeded");
+ if (!chunk_done.is_valid()) {
+ LOG_ERROR("batch_paged_attention_three_layer: chunk produced no completion task");
return;
}Also applies to: 189-193
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In
`@tests/st/a2a3/host_build_graph/batch_paged_attention_three_layer/kernels/orchestration/paged_attention_three_layer_orch.cpp`
around lines 83 - 90, Move the layer barrier capacity validation in the
orchestration routine so it runs before any tasks are submitted for a layer.
Reject the layer when q_loop * num_chunks exceeds MAX_LAYER_DONE_DEPS, returning
through the existing failure path before creating partial dependencies. Preserve
normal submission and barrier/final-epoch behavior for valid layer sizes.
| orch.submit_next_level(callables.paged_attention, chip_args, cfg, worker=0) | ||
| if request_id == _REQUEST_A: | ||
| request_a_submitted.set() | ||
| self._relay_tokens(queue, request_id, expected_epochs, emitter) | ||
|
|
||
| def consume(stream, output, errors): | ||
| try: | ||
| output.extend(stream) | ||
| except BaseException as exc: # noqa: BLE001 | ||
| errors.append(exc) | ||
|
|
||
| with worker.open_request_session(request_orch, max_active_runs=2) as session: | ||
| streams = { | ||
| _REQUEST_A: session.submit(requests[_REQUEST_A], config=config, request_id=_REQUEST_A), | ||
| } | ||
| assert request_a_submitted.wait(_STREAM_TIMEOUT_S) | ||
| # Place B's Host O behind A's first publication instead of profiling | ||
| # two simultaneous cold starts. | ||
| time.sleep(_CROSS_REQUEST_STAGGER_S) | ||
| streams[_REQUEST_B] = session.submit(requests[_REQUEST_B], config=config, request_id=_REQUEST_B) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Synchronize request B on A’s first token, not its submission.
request_a_submitted is set before _relay_tokens receives anything, so the sleep does not guarantee the stated “behind A’s first publication” ordering. Slow startup can profile two cold starts and make the overlap result nondeterministic. Set the event when sequence 1 is received, then submit B.
Proposed synchronization
-def _relay_tokens(self, queue, request_id: int, expected_epochs: int, emitter) -> None:
+def _relay_tokens(self, queue, request_id: int, expected_epochs: int, emitter, first_token_event=None) -> None:
for seq in range(1, expected_epochs + 1):
token, message = _read_token(queue, request_id, seq, expected_epochs)
+ if seq == 1 and first_token_event is not None:
+ first_token_event.set()
...
- if request_id == _REQUEST_A:
- request_a_submitted.set()
- self._relay_tokens(queue, request_id, expected_epochs, emitter)
+ self._relay_tokens(
+ queue,
+ request_id,
+ expected_epochs,
+ emitter,
+ request_a_submitted if request_id == _REQUEST_A else None,
+ )📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| orch.submit_next_level(callables.paged_attention, chip_args, cfg, worker=0) | |
| if request_id == _REQUEST_A: | |
| request_a_submitted.set() | |
| self._relay_tokens(queue, request_id, expected_epochs, emitter) | |
| def consume(stream, output, errors): | |
| try: | |
| output.extend(stream) | |
| except BaseException as exc: # noqa: BLE001 | |
| errors.append(exc) | |
| with worker.open_request_session(request_orch, max_active_runs=2) as session: | |
| streams = { | |
| _REQUEST_A: session.submit(requests[_REQUEST_A], config=config, request_id=_REQUEST_A), | |
| } | |
| assert request_a_submitted.wait(_STREAM_TIMEOUT_S) | |
| # Place B's Host O behind A's first publication instead of profiling | |
| # two simultaneous cold starts. | |
| time.sleep(_CROSS_REQUEST_STAGGER_S) | |
| streams[_REQUEST_B] = session.submit(requests[_REQUEST_B], config=config, request_id=_REQUEST_B) | |
| orch.submit_next_level(callables.paged_attention, chip_args, cfg, worker=0) | |
| self._relay_tokens( | |
| queue, | |
| request_id, | |
| expected_epochs, | |
| emitter, | |
| request_a_submitted if request_id == _REQUEST_A else None, | |
| ) | |
| def consume(stream, output, errors): | |
| try: | |
| output.extend(stream) | |
| except BaseException as exc: # noqa: BLE001 | |
| errors.append(exc) | |
| with worker.open_request_session(request_orch, max_active_runs=2) as session: | |
| streams = { | |
| _REQUEST_A: session.submit(requests[_REQUEST_A], config=config, request_id=_REQUEST_A), | |
| } | |
| assert request_a_submitted.wait(_STREAM_TIMEOUT_S) | |
| # Place B's Host O behind A's first publication instead of profiling | |
| # two simultaneous cold starts. | |
| time.sleep(_CROSS_REQUEST_STAGGER_S) | |
| streams[_REQUEST_B] = session.submit(requests[_REQUEST_B], config=config, request_id=_REQUEST_B) |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In
`@tests/st/a2a3/host_build_graph/batch_paged_attention_three_layer/test_batch_paged_attention_three_layer_l3.py`
around lines 223 - 242, Update the request synchronization around _relay_tokens
and request_a_submitted so the event is set only when A’s first token (sequence
1) is received and published, rather than immediately after submit. Keep B’s
session.submit after waiting for that event, preserving the existing stagger
only after A’s first publication.
| def test_emit_waits_until_consumer_receives_token() -> None: | ||
| stream = RequestStream(7, cast(RequestSession, _StreamSession())) | ||
| emitter = RequestEmitter(stream) | ||
| producer_done = threading.Event() | ||
|
|
||
| def produce() -> None: | ||
| emitter.emit("token-1") | ||
| producer_done.set() | ||
|
|
||
| thread = threading.Thread(target=produce) | ||
| thread.start() | ||
| time.sleep(0.02) | ||
| assert not producer_done.is_set() | ||
| assert stream.next(timeout=1.0) == "token-1" | ||
| assert producer_done.wait(1.0) | ||
| stream._finish(None) | ||
| thread.join() |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Synchronize producer entry before asserting that emit blocks.
The producer may not have started during the fixed sleep, so this test can pass without exercising the handoff. Signal immediately before emit() and wait for that signal before checking producer_done.
Proposed test fix
emitter = RequestEmitter(stream)
+ producer_started = threading.Event()
producer_done = threading.Event()
def produce() -> None:
+ producer_started.set()
emitter.emit("token-1")
producer_done.set()
thread = threading.Thread(target=produce)
thread.start()
- time.sleep(0.02)
+ assert producer_started.wait(1.0)
assert not producer_done.is_set()📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| def test_emit_waits_until_consumer_receives_token() -> None: | |
| stream = RequestStream(7, cast(RequestSession, _StreamSession())) | |
| emitter = RequestEmitter(stream) | |
| producer_done = threading.Event() | |
| def produce() -> None: | |
| emitter.emit("token-1") | |
| producer_done.set() | |
| thread = threading.Thread(target=produce) | |
| thread.start() | |
| time.sleep(0.02) | |
| assert not producer_done.is_set() | |
| assert stream.next(timeout=1.0) == "token-1" | |
| assert producer_done.wait(1.0) | |
| stream._finish(None) | |
| thread.join() | |
| def test_emit_waits_until_consumer_receives_token() -> None: | |
| stream = RequestStream(7, cast(RequestSession, _StreamSession())) | |
| emitter = RequestEmitter(stream) | |
| producer_started = threading.Event() | |
| producer_done = threading.Event() | |
| def produce() -> None: | |
| producer_started.set() | |
| emitter.emit("token-1") | |
| producer_done.set() | |
| thread = threading.Thread(target=produce) | |
| thread.start() | |
| assert producer_started.wait(1.0) | |
| assert not producer_done.is_set() | |
| assert stream.next(timeout=1.0) == "token-1" | |
| assert producer_done.wait(1.0) | |
| stream._finish(None) | |
| thread.join() |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tests/ut/py/test_worker/test_request_session.py` around lines 69 - 85, Update
test_emit_waits_until_consumer_receives_token to add a producer-entry
synchronization event set immediately before emitter.emit("token-1"), then wait
for that event before asserting producer_done is unset. Preserve the existing
token consumption, completion, stream cleanup, and thread-join assertions.
| def parse_log(log: Path) -> tuple[dict[int, list[dict]], list[dict], list[dict], list[int]]: | ||
| spans: dict[int, list[dict]] = {} | ||
| tokens: list[dict] = [] | ||
| deliveries: list[dict] = [] | ||
| for line in log.read_text(encoding="utf-8", errors="replace").splitlines(): | ||
| if match := STRACE_RE.search(line): | ||
| spans.setdefault(int(match.group("inv")), []).append( | ||
| { | ||
| "name": match.group("name"), | ||
| "ts": int(match.group("ts")), | ||
| "dur": int(match.group("dur")), | ||
| "attrs": parse_attrs(match.group("attrs")), | ||
| } | ||
| ) | ||
| if match := TOKEN_RE.search(line): | ||
| tokens.append( | ||
| { | ||
| "request": int(match.group("request")), | ||
| "seq": int(match.group("seq")), | ||
| "ts": int(match.group("ts")), | ||
| } | ||
| ) | ||
| if match := DELIVERY_RE.search(line): | ||
| deliveries.append( | ||
| { | ||
| "request": int(match.group("request")), | ||
| "seq": int(match.group("seq")), | ||
| "ts": int(match.group("ts")), | ||
| } | ||
| ) | ||
| request_ids = list(dict.fromkeys(token["request"] for token in tokens)) | ||
| if not spans or len(request_ids) < 2: | ||
| raise ValueError("log must contain STRACE spans and tokens for at least two requests") | ||
| return spans, tokens, deliveries, request_ids | ||
|
|
||
|
|
||
| def map_host_invocations(spans: dict[int, list[dict]], request_ids: list[int]) -> dict[int, int]: | ||
| normal_invs = [inv for inv, values in spans.items() if any(span["name"] == "simpler_run" for span in values)] | ||
| prepare_invs = [ | ||
| inv for inv, values in spans.items() if any(span["name"] == "simpler_prepare_request" for span in values) | ||
| ] | ||
| prepare_invs.sort( | ||
| key=lambda inv: next(span["ts"] for span in spans[inv] if span["name"] == "simpler_prepare_request") | ||
| ) | ||
| if len(normal_invs) != 1 or len(prepare_invs) != len(request_ids) - 1: | ||
| raise ValueError("log must contain one normal run and one Host prepare invocation per later request") | ||
| return dict(zip(request_ids, [normal_invs[0], *prepare_invs])) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Do not infer request-to-invocation mapping from token arrival order.
request_ids preserves first token-receipt order, which is nondeterministic under concurrent execution. If B emits first, it is incorrectly assigned A’s normal invocation and every O/S window is swapped. Parse the HOST_REQUEST_PREPARED request_id=... records and map prepared invocations explicitly; the remaining request owns the normal invocation.
🧰 Tools
🪛 Ruff (0.15.21)
[warning] 131-131: zip() without an explicit strict= parameter
Add explicit value for parameter strict=
(B905)
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tools/cross_request_host_device_perfetto.py` around lines 85 - 131, Update
parse_log and map_host_invocations so request-to-invocation mapping does not
depend on token arrival order. Parse HOST_REQUEST_PREPARED records and associate
each prepared request ID with its corresponding simpler_prepare_request
invocation, then assign the remaining request to the sole simpler_run
invocation. Preserve validation that exactly one normal invocation and one
prepare invocation exist for each later request, and reject missing or ambiguous
mappings.
| received = {(token["request"], token["seq"]): token["ts"] for token in tokens} | ||
| delivered = {(token["request"], token["seq"]): token["ts"] for token in deliveries} | ||
| shared_keys = received.keys() & delivered.keys() | ||
| delivery_ms = {key: (delivered[key] - received[key]) / 1_000_000.0 for key in shared_keys} | ||
| one_to_one = received.keys() == delivered.keys() and all(delay >= 0 for delay in delivery_ms.values()) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Count records before claiming strict one-to-one delivery.
These dictionaries collapse duplicate receipts or deliveries by (request, seq), so duplicate events can still make strict_one_to_one true. Compare Counter values first, then use the timestamp maps for latency calculations.
Proposed duplicate validation
+from collections import Counter
...
+ received_counts = Counter((token["request"], token["seq"]) for token in tokens)
+ delivered_counts = Counter((token["request"], token["seq"]) for token in deliveries)
received = {(token["request"], token["seq"]): token["ts"] for token in tokens}
delivered = {(token["request"], token["seq"]): token["ts"] for token in deliveries}
...
- one_to_one = received.keys() == delivered.keys() and all(delay >= 0 for delay in delivery_ms.values())
+ one_to_one = (
+ received_counts == delivered_counts
+ and all(count == 1 for count in received_counts.values())
+ and all(delay >= 0 for delay in delivery_ms.values())
+ )📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| received = {(token["request"], token["seq"]): token["ts"] for token in tokens} | |
| delivered = {(token["request"], token["seq"]): token["ts"] for token in deliveries} | |
| shared_keys = received.keys() & delivered.keys() | |
| delivery_ms = {key: (delivered[key] - received[key]) / 1_000_000.0 for key in shared_keys} | |
| one_to_one = received.keys() == delivered.keys() and all(delay >= 0 for delay in delivery_ms.values()) | |
| received_counts = Counter((token["request"], token["seq"]) for token in tokens) | |
| delivered_counts = Counter((token["request"], token["seq"]) for token in deliveries) | |
| received = {(token["request"], token["seq"]): token["ts"] for token in tokens} | |
| delivered = {(token["request"], token["seq"]): token["ts"] for token in deliveries} | |
| shared_keys = received.keys() & delivered.keys() | |
| delivery_ms = {key: (delivered[key] - received[key]) / 1_000_000.0 for key in shared_keys} | |
| one_to_one = ( | |
| received_counts == delivered_counts | |
| and all(count == 1 for count in received_counts.values()) | |
| and all(delay >= 0 for delay in delivery_ms.values()) | |
| ) |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tools/cross_request_host_device_perfetto.py` around lines 187 - 191, Update
the delivery validation around received, delivered, and one_to_one to count
every (request, seq) occurrence with Counter before declaring strict one-to-one
delivery. Require the Counter values to match, while retaining the existing
timestamp-map calculation for shared-key latency and nonnegative delays.
7e31f63 to
34d344c
Compare
Review 后的架构层反馈:建议先对齐方案,再拆分先说结论:这个 PR 能达成"O/S 跨 token 流水线 + 跨请求重叠"的目标,设计很用心。但核心并发机制偏复杂(独立控制通道 admission lane + 显式 1. 一个关键观察:当前方案只拿到 O∥S,没拿到 S∥S
但 也就是说,当前这套 admission lane + 两阶段 API,换来的只是 2. 一个更简洁的替代方案:常驻消息线程 + run 派工作线程把主循环改成"收到
代价集中在一处:mailbox 从"单槽同步"升级为"多槽 / 异步完成"(派线程后不能立即写 需要权衡的一点:显式 3. 一个统一模型:bank 跟着线程走
唯一需要的小改动: 4. 两个必须先修的死锁(与并发触发方式正交,走哪个方案都要修)
二者当前都没有回归测试( 5. 关于拆分方案对齐后再拆分会自然很多。core 3300+ 行、跨四个可分离机制(arena banks / prepare-execute ABI / epoch 协议 / 并发 Orchestrator+RequestSession)。建议先合"行为等价的底座"(HostApi 线程上下文、scheduler 建议先就以上(尤其第 2 点 mailbox 是否多槽化)写一份简短 design-note 对齐,再决定拆分边界。 |
34d344c to
d94da9b
Compare
|
@ChaoWao The branch has been rebased onto the current upstream main and updated with the architecture-aligned implementation discussed in the review. Architecture update:
Rebase compatibility fix:
Validation on a2a3 hardware:
No simulator run was performed in this iteration, per the agreed scope; GitHub sim CI remains the final automatic regression step. The final Python file passes all applicable pre-commit hooks. The broader clang-tidy hook is still blocked before analysis by a stale absolute simpler_setup path in the cluster environment; the other relevant hooks pass. |
79c286f to
6733422
Compare
|
CI follow-up for the current head
Validation:
The first macOS packaging retry encountered a runner-level |
6733422 to
81675ce
Compare
81675ce to
8509dec
Compare
|
Rebased this PR onto Architecture update
Latest-main compatibilityThe rebase includes the polling-completion scheduler and removal of legacy graph-output metadata from current
Validation on commit
|
|
CI follow-up: the first automatic The exact formatter output is applied in GitHub currently reports an Actions partial outage and degraded Webhooks/Issues service, so PR 1379's displayed head is temporarily lagging behind the fork branch. I am leaving the GitHub sim suite to rerun automatically when the synchronize event is processed. |
a3c2b2a to
95f2582
Compare
|
Fixed the remaining macOS Root cause: the two contract-driven mailbox slots can enter the same a2a3 sim The fix adds the missing per-runner Gate B to a2a3 sim. Host O/binding remains outside this lock, so the intended O||S pipeline is preserved; only the shared Device S runner state is serialized, matching onboard semantics. Validation on exact head
Per the current validation plan, I did not run sim manually; the refreshed GitHub sim jobs are the automatic regression check for this fix. |
- publish host graph epochs for overlapped device execution and preserve a final epoch for ordinary orchestrations - add host request admission, prepared request execution, and concurrent Worker run support - integrate current L3-L2 direct-region behavior and add tracing, documentation, and hardware coverage - fix CI formatting, lint, type checking, copyright, and current-main compatibility issues
22dc243 to
7fd379f
Compare
|
CI follow-up after rebasing onto current
Validation on commit
The branch was rebased and force-updated with lease. A fresh GitHub Actions run should now validate the merge result. |
|
Follow-up fix for the remaining macOS
Validation at exact commit
The prior A5 failures appear runner-infrastructure-related ( |
|
CI follow-up for run 30006827972:
The remaining red jobs are self-hosted hardware failures:
The current token receives HTTP 403 when rerunning individual jobs. A maintainer needs to restore the A5 device allocation and rerun these three self-hosted jobs. |
3 participants
Summary
This PR enables host-side orchestration and device-side scheduling to run as a pipeline.
Worker.run()calls. While one request is running, another request can be accepted and prepared.Profiling Results
The dual-request trace demonstrates:
Worker.run()calls.