perf(deepseek): overlap decode prepare and reclaim

[high-cloud]

Summary

• pipeline decode preparation through two reusable input slots while the prior device invocation runs
• make prepare produce a complete fused-MTP dispatch snapshot, leaving only the blocking PyPTO runtime call on the device lane
• reclaim ping-ponged outputs on an independent FIFO output lane while the next prepared decode dispatches
• keep per-request MTP tail, draft, generation, and committed state in stable device slots
• remove the obsolete deferred arm/launch API and its armed-ticket worker branch
• remove request-layout assignment caching, static-metadata cache keys, cross-step L3 argument-tuple caches, and their invalidation lifecycle
• rebuild request-specific metadata and dispatch arguments normally on the prepare lane; retain only correctness-required ping-pong slots and persistent MTP device state
• finalize MTP prefill state inside the prefill command and preserve synchronous fallbacks for unsupported paths
• add profiling spans and protocol coverage for preparation, independent reclaim, lifecycle release, and tail flushing
• harden slot ownership, lifecycle ordering, malformed-command FIFO handling, and host-embedding fallback following review

Why

The asynchronous scheduler could queue future decode steps, but host preparation and output processing still serialized the device lane. Persistent device state removes the prior-token dependency, and independent FIFO reclaim allows the next prepared invocation to launch without waiting for scheduler-visible output processing.

PR #139 reduced serial host gaps by memoizing request assignments, static metadata, and complete L3 argument tuples. Those caches were tied to stable request/rank layouts, had narrow hit conditions, and required extra invalidation logic. Once prepare runs concurrently with the much longer device invocation, that work is no longer on the critical path. This PR therefore prefers normal per-step construction on the prepare lane over topology- and lifecycle-sensitive caches.

Performance: before vs after this PR

The baseline is the final synchronous decode path from PR #139 (a74a514, merged as b1a35ab). The optimized result is this PR after removing the temporary cache shortcuts (1f7b270, rebased without semantic changes as a3a19e0). Both profiles used the same model, prompt, exactly 20 output tokens, fused one-L2 MTP decode, PTOAS 0.54, host STRACE with device STRACE disabled, and had MTP acceptance 10/11 (90.91%).

Steady decode metric	Before this PR	After this PR	Change
Serving fused-decode span	47.206 ms	45.425 ms	-3.8%
Critical-rank host STRACE	44.785 ms	42.752 ms	-4.5%
Cross-round host dispatch gap	4.426 ms	0.375 ms	-91.5%
Decode span + cross-round gap	51.632 ms	45.800 ms	-11.3%
Prepare(N+1) overlaps dispatch(N)	0/10	10/10	fully overlapped

The key framework result is the 91.5% reduction in the exposed cross-round host gap. All ten successor prepares completed before the preceding device invocation ended; the minimum remaining device-execution margin was 27.054 ms. Reclaim ran on the independent output lane and no longer blocked the next device dispatch.

The baseline was task task_20260805_024605_15915111219 on devices 8-15. The optimized run was task task_20260810_034449_263105568 on devices 0,2,4,6,8,10,12,14. Because the physical device sets differ, the 3.8%-4.5% decode-duration changes are directional; the lane ordering and exposed-gap reduction are the primary evidence for this PR. HTTP wall time is intentionally excluded because prefill compilation/runtime variance dominated it.

Baseline artifact: artifacts/dsv4-profile-hostopt-rebased-p1c66-r3165-lc46d-ptoas054-20260805 in the PR #139 checkout.

Optimized artifact: artifacts/dsv4-serving-profile-20260810-even-cleanup.

Current validation

• Current commit: a109161 on origin/main@0033990
• python -m pytest tests/unit -q — 139 passed after review fixes
• targeted Ruff checks on all changed Python files
• header, English-only, Python compile, and diff checks
• optimized real-NPU profile exited 0 with exactly 20 output tokens
• 11 fused decode, 11 prepare, 11 dispatch, and 11 reclaim spans; zero arm spans
• eight host STRACE lanes captured; device STRACE intentionally disabled, so Device Effective is unavailable

The NPU profile was captured on the pre-rebase equivalent implementation (1f7b270). The rebase resolved the new upstream KernelCompiler integration and worker execution conflict; the current commit has host validation but has not been rerun on NPU.

Review hardening

• each of the two decode slots remains owned until output reclaim has finished reading its captured tensors
• request release executes on the FIFO device lane, with cache-entry identity protection for later same-ID registrations
• speculative MTP requests wait for confirmed terminal prefill before first-decode scheduling, without adding steady device-lane work
• the obsolete no-op executor session hook and its worker call sites were removed
• malformed commands travel through the ordered device/output lanes rather than bypassing earlier results
• executors without device decode embedding fall back to synchronous preparation instead of embedding placeholder token zero
• negative slot indices are rejected, active-slot sampled outputs are used, MTP state descriptor reads are locked, and prepared buffer sizing uses the model vocabulary

Earlier multi-request validation

Four sequential requests with 20, 17, 13, and 9 output tokens completed, and two staggered four-request waves completed all eight requested outputs. In that trace, 34/34 steady transitions overlapped prepare(N+1) with dispatch(N) and reclaim(N) with dispatch(N+1).

The two waves used different packed-prefill groupings (1+3 and 3+1). Greedy output matched between waves for 2/4 prompts; two prompts differed despite matching output-token counts. That batch-shape-dependent precision observation remains a follow-up and is not claimed as a precision pass for the current commit.

Depends on hw-native-sys/pypto-lib#917.

[coderabbitai]

📝 Walkthrough

Walkthrough

The serving worker now supports asynchronous prepared decode with split device dispatch and output reclamation. DeepSeek V4 uses ping-pong buffers and persistent generation-aware MTP device state. Prefill finalization initializes state, and tests cover ordering, reuse, EOS handling, and reclamation.

Changes

Asynchronous MTP decode

Layer / File(s)	Summary
Executor lifecycle contracts pypto_serving/model/common/executor/executor.py, pypto_serving/model/common/executor/pypto_executor.py	Executors now expose hooks for prefill finalization, prepared decode, split dispatch and reclaim, and token dependency checks. Unsupported runners use fallback behavior.
Worker pipeline orchestration pypto_serving/serving/server/serving_worker.py, tests/unit/serving/engine/test_async_pipeline.py, tests/unit/serving/server/test_worker_step_protocol.py, tests/unit/serving/device_sampling_fakes.py	The worker now prepares decode work asynchronously, preserves FIFO ordering, reclaims results in order, handles EOS and shutdown, and finalizes terminal prefill state.
DeepSeek persistent MTP state pypto_serving/model/deepseek/npu_runner.py, pypto_serving/model/deepseek/npu_executor.py, docs/dev/model/deepseek-v4.md, pypto-lib, .agents/skills/profile-dsv4-serving-strace/scripts/analyze_profile.py	DeepSeek V4 now uses isolated decode slots, persistent MTP state tensors, generation-tagged slot reuse, prepared fused arguments, and split dispatch and reclaim. Supporting documentation, kernel imports, profiling compatibility, and the pypto-lib reference were updated.
DeepSeek prepared decode validation tests/unit/model/deepseek/test_model_components.py	Tests cover ping-pong slot isolation, metadata rebuilding, state reservation, generation routing, prepared dispatch, and device-owned recurrent state.

Estimated code review effort: 5 (Critical) | ~120 minutes

Possibly related PRs

• hw-native-sys/pypto-serving#89: Earlier DeepSeek V4 MTP implementation extended by persistent device state and asynchronous execution.
• hw-native-sys/pypto-serving#137: Related asynchronous pipeline, worker protocol, and DeepSeek component test changes.
• hw-native-sys/pypto-serving#145: Related worker, executor, and DeepSeek MTP changes for fixed-shape decode.

Sequence Diagram(s)

sequenceDiagram
  participant ServingWorker
  participant PyptoExecutor
  participant DeepSeekV4ModelRunner
  participant DeviceKernel
  participant ResultReclaimer

  ServingWorker->>PyptoExecutor: Prepare decode snapshot
  PyptoExecutor->>DeepSeekV4ModelRunner: Prepare slot-local inputs
  ServingWorker->>PyptoExecutor: Dispatch prepared decode
  PyptoExecutor->>DeepSeekV4ModelRunner: Launch device work
  DeepSeekV4ModelRunner->>DeviceKernel: Execute fused MTP decode
  DeviceKernel->>DeepSeekV4ModelRunner: Update persistent device state
  ServingWorker->>ResultReclaimer: Queue pending result
  ResultReclaimer->>PyptoExecutor: Reclaim prepared decode
  PyptoExecutor->>DeepSeekV4ModelRunner: Read committed outputs

Loading

Poem

I’m a rabbit in the decode lane,
Hopping through slots without stale-state pain.
Prefill plants tokens, kernels keep score,
FIFO results reach the waiting door.
Ping-pong buffers make the pipeline flow—
Async MTP is ready to go! 🐇

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 59.83% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Description check	✅ Passed	The description clearly explains the decode overlap, MTP state changes, fallback behavior, validation, and performance results covered by the changeset.
Title check	✅ Passed	The title clearly and concisely identifies the primary change: overlapping DeepSeek decode preparation and output reclamation.

✨ Finishing Touches 💡 1

⚔️ Resolve merge conflicts 💡

• Resolve merge conflict in branch agent/async-deepseek-decode-pipeline

---

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.}

[coderabbitai]

Actionable comments posted: 6

🧹 Nitpick comments (4)

pypto_serving/model/deepseek/npu_runner.py (4)

3602-3620: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Read _mtp_request_states under _mtp_state_lock.

Line 3608 uses self._mtp_request_states.get(request_id) without the lock. The sibling helper _require_mtp_request_state takes _mtp_state_lock for the same read, and release_finished_requests pops from the dict on the worker output lane while this method runs on the prepare lane.

The read itself is atomic, so the dict cannot be corrupted. Use the lock for consistency with the other accessors and to keep the state lookup and the tail_slot_id/generation reads in one critical section.

🤖 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 `@pypto_serving/model/deepseek/npu_runner.py` around lines 3602 - 3620, Protect
the MTP state lookup and related field reads in the loop around
`_mtp_request_states` with `_mtp_state_lock`. Keep the lookup, `tail_slot_id`,
and any `generation` access in the same critical section, matching
`_require_mtp_request_state`, while preserving the existing missing-state and
unreserved-tail behavior.

---

2636-2666: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Simplify the state collection in _reclaim_mtp_decode.

Line 2662 appends pending.states[len(states)], which depends on the current length of the list being built. The loop variable request_id is never used, which Ruff reports as B007 at line 2645.

Zip pending.states into the same loop and drop the local states list.

♻️ Proposed refactor

             accepted_counts_list = []
             accepted = []
-            states = []
-            for request_id, rank, local_row in zip(
-                inputs.request_ids,
+            states = pending.states
+            for rank, local_row in zip(
                 inputs.ranks,
                 inputs.local_rows,
                 strict=True,
             ):
                 row_start = local_row * decode_seq
                 main_tokens = pending.sampled_ids[
                     rank,
                     row_start : row_start + decode_seq,
                     0,
                 ].tolist()
                 accepted_count = int(
                     pending.accepted_counts[rank, local_row].item()
                 )
                 accepted_counts_list.append(accepted_count)
                 accepted.append([int(token) for token in main_tokens[:accepted_count]])
-                states.append(pending.states[len(states)])

🤖 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 `@pypto_serving/model/deepseek/npu_runner.py` around lines 2636 - 2666, Update
_reclaim_mtp_decode to iterate over pending.states alongside inputs.request_ids,
inputs.ranks, and inputs.local_rows in the existing strict zip, append accepted
tokens using the provided state directly, and remove the temporary states list
and pending.states[len(states)] lookup. Rename the unused request_id loop
variable appropriately so Ruff no longer reports B007.

Source: Linters/SAST tools

---

3777-3785: 🗄️ Data Integrity & Integration | 🔵 Trivial | ⚡ Quick win

_DeepSeekV4DecodeSharedBuffers pins sampled_ids and tensors to slot 0.

Both fields reference self._decode_input_slots[0]. _execute_main_decode uses decode_buffers.sampled_ids, so the autoregressive prepared path writes slot 0 even when prepared.buffer_slot is 1.

This is safe today because supports_async_decode_reclaim returns self._compiled.enable_mtp, so the non-MTP path never splits dispatch and reclaim and the device lane serializes dispatch with the result read. The coupling is implicit and breaks if split reclaim is extended to the non-MTP path.

Select the active slot in _execute_main_decode from prepared.buffer_slot, or document why slot 0 is correct there.

🤖 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 `@pypto_serving/model/deepseek/npu_runner.py` around lines 3777 - 3785, The
_DeepSeekV4DecodeSharedBuffers construction hardcodes sampled_ids and tensors
from decode input slot 0, while _execute_main_decode may operate on another
prepared slot. Select the active decode input slot using prepared.buffer_slot
when constructing or consuming these fields, preserving the existing buffer-slot
selection for the autoregressive prepared path.

---

3109-3130: 🎯 Functional Correctness | 🔵 Trivial | 💤 Low value

Remove the duplicate DeepSeek V4 vocab-size constant.

_require_decode_logits_buffer() returns the cached allocation without checking the requested size. _require_decode_logits_buffer(DEEPSEEK_V4_VOCAB_SIZE) and _require_decode_logits_buffer(model.config.vocab_size) currently use the same constant, but keeping both makes the size assumption less clear. Use model.config.vocab_size here with the same validation that load_lm_head_weights() applies.

🤖 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 `@pypto_serving/model/deepseek/npu_runner.py` around lines 3109 - 3130, Update
_bind_prepared_mtp_dispatch to request the logits buffer using
model.config.vocab_size instead of DEEPSEEK_V4_VOCAB_SIZE, matching the
validation and size source used by load_lm_head_weights(). Keep the existing
buffer retrieval and dispatch argument construction unchanged.

🤖 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 `@pypto_serving/model/deepseek/npu_runner.py`:
- Around line 2179-2183: Validate buffer_slot explicitly in the decode buffer
access near _ensure_decode_buffers, rejecting values below 0 as well as values
above the available slot range before indexing _decode_input_slots. Preserve the
existing ValueError message/handling for invalid slots and the normal staged
lookup for slots 0 and 1.
- Around line 3663-3786: Update _decode_input_slots and the step-slot ownership
in WorkerProcess._pipelined_busy_loop so the number of reusable decode slots
matches the maximum in-flight device steps. Ensure dispatch does not replace or
recycle a slot until reclaim has finished reading all captured input and output
tensors, rather than relying on cmd.step_id % 2 when two-step overlap can reuse
a slot.

In `@pypto_serving/serving/server/serving_worker.py`:
- Around line 277-297: Update
pypto_serving/serving/server/serving_worker.py:277-297 around the StepCommand
registration and _reclaim_pending_decode lifecycle to remove the outdated
FIFO-ordering claim and ensure reclaiming step N cannot remove the _req_cache
entry registered for step N+1, using step/generation identity validation. Update
pypto_serving/model/deepseek/npu_runner.py:1696-1706 to correct the
FIFO-device-lane comment and verify that the relevant kernel rejects descriptors
whose generation no longer matches the recycled slot, enforcing that validation
if necessary.
- Around line 390-408: Update _reclaim_pending_decode so
_release_finished_request_state(work.cmd.finished_request_ids) executes in a
finally block, regardless of whether reclaim_prepared_decode or
_consume_decode_result raises. Preserve the existing successful StepResult path
and exception logging/error response while ensuring finished request state is
always released.
- Around line 269-276: Update the decode-exception path in
_device_execution_loop to stop writing directly to self.output_queue; wrap the
encoded StepResult error in the marker type _CompletedStepOutput and enqueue it
through output_work_queue so the output lane preserves FIFO ordering with
earlier steps. Ensure the output lane forwards the completed failure result
using the existing ordered-output handling.
- Around line 678-709: Prevent _prepare_step_command from preparing decode work
when the executor supports async decode preparation but not device decode
embeddings, since the placeholder batch would bind hidden_states to token zero.
Add an explicit validation failure before _make_decode_batch, or alternatively
update _late_bind_prepared_decode_batch to re-bind hidden_states whenever
token_ids are replaced; preserve correct embeddings for all prepared decode
batches.

---

Nitpick comments:
In `@pypto_serving/model/deepseek/npu_runner.py`:
- Around line 3602-3620: Protect the MTP state lookup and related field reads in
the loop around `_mtp_request_states` with `_mtp_state_lock`. Keep the lookup,
`tail_slot_id`, and any `generation` access in the same critical section,
matching `_require_mtp_request_state`, while preserving the existing
missing-state and unreserved-tail behavior.
- Around line 2636-2666: Update _reclaim_mtp_decode to iterate over
pending.states alongside inputs.request_ids, inputs.ranks, and inputs.local_rows
in the existing strict zip, append accepted tokens using the provided state
directly, and remove the temporary states list and pending.states[len(states)]
lookup. Rename the unused request_id loop variable appropriately so Ruff no
longer reports B007.
- Around line 3777-3785: The _DeepSeekV4DecodeSharedBuffers construction
hardcodes sampled_ids and tensors from decode input slot 0, while
_execute_main_decode may operate on another prepared slot. Select the active
decode input slot using prepared.buffer_slot when constructing or consuming
these fields, preserving the existing buffer-slot selection for the
autoregressive prepared path.
- Around line 3109-3130: Update _bind_prepared_mtp_dispatch to request the
logits buffer using model.config.vocab_size instead of DEEPSEEK_V4_VOCAB_SIZE,
matching the validation and size source used by load_lm_head_weights(). Keep the
existing buffer retrieval and dispatch argument construction unchanged.

🪄 Autofix

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 Plus

Run ID: 8c3bc95f-dda7-4a96-bb66-a2215f845ea1

📥 Commits

Reviewing files that changed from the base of the PR and between a181044 and 1f7b270.

📒 Files selected for processing (12)

• .agents/skills/profile-dsv4-serving-strace/scripts/analyze_profile.py
• docs/dev/model/deepseek-v4.md
• pypto-lib
• pypto_serving/model/common/executor/executor.py
• pypto_serving/model/common/executor/pypto_executor.py
• pypto_serving/model/deepseek/npu_executor.py
• pypto_serving/model/deepseek/npu_runner.py
• pypto_serving/serving/server/serving_worker.py
• tests/unit/model/deepseek/test_model_components.py
• tests/unit/serving/device_sampling_fakes.py
• tests/unit/serving/engine/test_async_pipeline.py
• tests/unit/serving/server/test_worker_step_protocol.py