feat(deepseek): expand continuous batching across fixed kernels

[hashiqiqixian]

Summary - expand DeepSeek V4 prefill from one request per DP rank to four partition-local slots, while retaining compiler-compatible B1S128 kernel calls - allow scheduler steps to mix decode and prefill work, with decode dispatched first for latency - split decode work beyond one fixed launch into partition-safe micro-batches and merge their outputs in one worker step - preserve request-local MTP committed-tail state across micro-batches by paging it through the fixed B4 device pool - remove the serving-wide single-dispatch request cap and document the resulting fixed-shape concurrency model ## Scope and ownership This PR changes the pypto-serving scheduler, worker dispatch, DeepSeek V4 runner, configuration, tests, and developer documentation. It does not change pypto-lib or the compiled kernel ABI. The fixed launch specializations remain: - prefill: up to four B1S128 calls per DP rank in one worker batch - MTP decode: B4S2 per rank, K=1, up to 32 requests per launch across 8 DP ranks - non-MTP decode: B8S1 per rank, up to 64 requests per launch across 8 DP ranks Larger active sets are served by additional worker dispatches. This increases functional serving concurrency, not per-launch kernel throughput, and extra launches plus MTP tail paging add latency. ## Out of scope - MTP K greater than 1 - dynamic-shape or unified prefill/decode kernels - changing the B4S2/B8S1 kernel specializations - pypto-lib kernel or ABI changes - hardware performance claims or benchmark targets ## Validation

Current validated revisions:

• pypto-serving@ee5f21f0d9efb63db844c115f4a9712c289c0169
• pypto-lib@f0d352ea2a40cefca8aa26873261d6e19e0eefec
• pypto@1c66f01c180e1fc3a858817dfb8f52d6ecefff88

CPU and static checks:

• remote focused unit tests for the DeepSeek runner, offline path, scheduler, and worker protocol: 78 passed
• new staging regressions (shared prefill buffer reuse, per-slot argument caching, and one tail transfer per active rank): 3 passed
• Ruff on the changed runner and test file: passed
• local python -m py_compile and git diff --check: passed

8-NPU DeepSeek V4 Flash W8A8 functional validation on devices 8-15:

• task: task_20260809_021107_78789517123, exit 0
• configuration: DP8/EP8, legacy MTP K=1, max_num_seqs=16
• nine simultaneous HTTP completion requests: 9/9 passed
• every response returned the exact expected 10-token completion: a leading global information and communications technology (ICT)
• MTP acceptance reached 45 accepted out of 54 proposed tokens (83.33%), so both accept and reject paths executed
• all eight chip processes became ready; the server passed health, completed all requests, stopped, and released devices 8-15
• artifacts: /data/chenshenai/test1/pypto-serving/.validation/full-mtp-20260809_021107

Shutdown emitted one PyPTO runner close-budget error for a child process, followed by normal application shutdown. The task still exited 0, no serving/chip process remained, and all device locks were released. This validation is a functional concurrency result, not a throughput benchmark or performance claim.

[coderabbitai]

Important

Review skipped

Auto incremental reviews are disabled on this repository.

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 Plus

Run ID: 1b88e1fb-36d4-4435-b5c9-eb64ea70d52f

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

📝 Walkthrough

Walkthrough

DeepSeek V4 serving now supports four-slot rank-local prefill, capacity-aware decode and prefill dispatch, mixed decode/prefill scheduling, and request-local MTP tail-state paging. Runtime configuration and documentation expose the new limits.

Changes

DeepSeek V4 serving

Layer / File(s)	Summary
Runtime capacity contracts pypto_serving/config/types.py, pypto_serving/cli/main.py, pypto_serving/model/common/executor/executor.py, pypto_serving/model/deepseek/npu_executor.py, pypto_serving/model/deepseek/offline.py, tests/unit/model/deepseek/test_offline.py, tests/unit/model/deepseek/test_model_components.py	Runtime configuration now carries a per-request prefill limit. Executors expose global and partition-local capacities. DeepSeek V4 selects MTP-dependent decode layouts and accepts max_num_seqs=64.
Mixed scheduling and partition dispatch pypto_serving/serving/engine/async_engine.py, pypto_serving/serving/sched/scheduler.py, pypto_serving/serving/server/serving_worker.py, tests/unit/serving/sched/test_async_scheduler.py, tests/unit/serving/server/test_worker_step_protocol.py	Scheduling applies shared prefill limits and admits prefill during decode. Worker dispatch runs decode first and partitions both operations under global and partition-local capacities.
Rank-local multi-slot prefill pypto_serving/model/deepseek/npu_runner.py, tests/unit/model/deepseek/test_model_components.py	Prefill supports four local slots per rank. Inputs, metadata, buffers, outputs, and logits use rank/slot-major layouts. Active slots dispatch independently with capacity validation.
Request-local MTP state and serving contract pypto_serving/model/deepseek/npu_runner.py, docs/dev/model/deepseek-v4.md, tests/test_deepseek_v4_accuracy.py, tests/unit/model/deepseek/test_model_components.py	MTP tail hidden state uses request-local host storage and fixed device-pool paging. Documentation and tests describe fixed decode shapes, mixed admission, prefill limits, and the 128-token threshold.

Estimated code review effort: 4 (Complex) | ~60 minutes

Sequence Diagram(s)

sequenceDiagram
  participant Scheduler
  participant ReplicaEngineCore
  participant serving_worker
  participant DeepSeekV4ModelRunner
  Scheduler->>ReplicaEngineCore: select mixed decode and prefill work
  ReplicaEngineCore->>serving_worker: execute step
  serving_worker->>serving_worker: partition decode before prefill
  serving_worker->>DeepSeekV4ModelRunner: dispatch fixed-capacity batches
  DeepSeekV4ModelRunner->>DeepSeekV4ModelRunner: page MTP tail state and run slots
  DeepSeekV4ModelRunner-->>serving_worker: return logits and generated tokens

Loading

Possibly related issues

• [Bug] DeepSeek V4 chunked prefill does not honor the fixed 128-token kernel limit #123 — Covers the 128-token DeepSeek V4 prefill limit and scheduler enforcement.
• [Feature] Add DeepSeekV4 Flash serving support #62 — Covers multi-slot batching, capacity-aware scheduling, and MTP state handling.

Possibly related PRs

• hw-native-sys/pypto-serving#113 — Introduces DeepSeek cache-layout and executor abstractions extended by this PR.
• hw-native-sys/pypto-serving#120 — Provides related DeepSeek V4 MTP and runtime changes.
• hw-native-sys/pypto-serving#132 — Modifies the same 128-token prefill and MTP chunking paths.

Poem

I’m a rabbit in the rank-local run,
Four slots hop beneath the sun.
Decode goes first, then prefills flow,
MTP tails page in and go.
Fixed shapes keep each launch bright—
128 tokens, packed just right.

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 48.75% 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.
Title check	✅ Passed	The title clearly summarizes the main change: expanded DeepSeek continuous batching across fixed kernels.
Description check	✅ Passed	The description directly explains the batching, scheduling, worker, MTP state, configuration, testing, and validation changes.

---

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]

🧹 Nitpick comments (4)

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

4546-4570: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Consider sharing the local-row assignment logic with _decode_assignment.

_prefill_assignment and _decode_assignment have the same body: dense per-rank local-row assignment with a capacity guard. They differ only in the capacity field (prefill_batch versus decode_batch), the error text, and decode's assignment cache. A shared private helper that takes ranks, the capacity, and a label would remove the duplication.

The return type is also _DeepSeekV4DecodeAssignment, and the prefill caller at Line 1769 uses only local_rows. Renaming the dataclass to a mode-neutral name would make the shared use clear.

🤖 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 4546 - 4570, Extract
the duplicated dense per-rank local-row assignment from _prefill_assignment and
_decode_assignment into a shared private helper accepting ranks, capacity, and a
mode label, while preserving each method’s capacity validation and error
wording. Reuse the helper from both assignment paths, retain decode’s assignment
cache, and rename _DeepSeekV4DecodeAssignment to a mode-neutral type name with
all references updated.

---

2744-2782: 🚀 Performance & Scalability | 🔵 Trivial | 💤 Low value

Consider caching _prefill_fwd_args per slot.

The staged buffers are fixed shared tensors, so the argument tuple for a given slot is stable across prefill steps. This method now runs once per active slot, so it rebuilds the weight dict and re-slices 21 tensors up to four times per step. _decode_fwd_args already caches its tuple in _decode_fwd_args_cache. A dict[int, tuple[Any, ...]] keyed by slot would give the same benefit here.

The transposed (prefill_batch, ranks, ...) storage that makes [:, slot] contiguous is correct and worth keeping.

🤖 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 2744 - 2782, Cache
the tuple produced by the prefill argument builder per slot, following the
existing _decode_fwd_args_cache pattern. In the method containing
_PREFILL_FWD_RANK_SLOT_NAMES and _ordered_layer_args, return the cached tuple
for a valid slot, and store the newly built result under that slot after
constructing and resident-marking the values. Preserve the existing [:, slot]
slicing and contiguity checks.

---

908-934: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick win

Consider writing prefill_x_hc directly into the shared buffer.

prefill_batch grew from 1 to 4, so this path now allocates and fills ranks × prefill_batch × token_rows × hc_mult × hidden float32 twice per prefill step: once for the clone() at Line 915 and once for the contiguous() at Line 933. The result is then copied again into buffers.x_hc by _stage_prefill_fwd_inputs. With the production layout this is roughly a 4x increase over the previous shape.

_pack_decode_x_hc already supports an out= destination for exactly this reason. An equivalent out= parameter here would let prepare_prefill_inputs fill the shared prefill buffer in place and remove both temporaries.

🤖 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 908 - 934, Update
the prefill packing helper around `rank_rows` to accept an `out` destination,
write `prefill_x_hc` directly into it, and avoid both the intermediate `clone()`
and final `contiguous()` allocations. Update `prepare_prefill_inputs` and
`_stage_prefill_fwd_inputs` to pass and reuse the shared `buffers.x_hc` buffer,
preserving the existing rank/local-row placement and output layout.

---

3110-3166: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick win

Consider batching the tail paging transfers per rank.

Both methods issue one copy_to/copy_from per request. With the full decode set this is up to 32 small device transfers in each direction per decode step, on the request thread.

buffers.tail_init_hidden[rank] is contiguous over the whole (decode_batch, hc_mult, hidden) block, and the pool shard has the same layout. Staging all rows for a rank first and then issuing one transfer per rank would reduce this to ranks transfers per direction.

The tradeoff is that a whole-rank transfer also moves rows for inactive local slots. Confirm that paging those rows in and out is harmless before making the change.

🤖 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 3110 - 3166, Batch
MTP tail paging in _page_in_mtp_tail_hidden and _page_out_mtp_tail_hidden by
staging all request rows for each rank, then issuing one contiguous transfer per
rank instead of one transfer per request. Preserve the existing request-state
validation and copy updated rows back to their corresponding request state;
confirm that transferring inactive local slots within each rank is harmless
before applying whole-rank transfers.

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

Nitpick comments:
In `@pypto_serving/model/deepseek/npu_runner.py`:
- Around line 4546-4570: Extract the duplicated dense per-rank local-row
assignment from _prefill_assignment and _decode_assignment into a shared private
helper accepting ranks, capacity, and a mode label, while preserving each
method’s capacity validation and error wording. Reuse the helper from both
assignment paths, retain decode’s assignment cache, and rename
_DeepSeekV4DecodeAssignment to a mode-neutral type name with all references
updated.
- Around line 2744-2782: Cache the tuple produced by the prefill argument
builder per slot, following the existing _decode_fwd_args_cache pattern. In the
method containing _PREFILL_FWD_RANK_SLOT_NAMES and _ordered_layer_args, return
the cached tuple for a valid slot, and store the newly built result under that
slot after constructing and resident-marking the values. Preserve the existing
[:, slot] slicing and contiguity checks.
- Around line 908-934: Update the prefill packing helper around `rank_rows` to
accept an `out` destination, write `prefill_x_hc` directly into it, and avoid
both the intermediate `clone()` and final `contiguous()` allocations. Update
`prepare_prefill_inputs` and `_stage_prefill_fwd_inputs` to pass and reuse the
shared `buffers.x_hc` buffer, preserving the existing rank/local-row placement
and output layout.
- Around line 3110-3166: Batch MTP tail paging in _page_in_mtp_tail_hidden and
_page_out_mtp_tail_hidden by staging all request rows for each rank, then
issuing one contiguous transfer per rank instead of one transfer per request.
Preserve the existing request-state validation and copy updated rows back to
their corresponding request state; confirm that transferring inactive local
slots within each rank is harmless before applying whole-rank transfers.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro Plus

Run ID: 545e5608-d9fd-4c19-aa83-0c4f65b2cb87

📥 Commits

Reviewing files that changed from the base of the PR and between 272b874 and d46a424.

📒 Files selected for processing (15)

• docs/dev/model/deepseek-v4.md
• pypto_serving/cli/main.py
• pypto_serving/config/types.py
• pypto_serving/model/common/executor/executor.py
• pypto_serving/model/deepseek/npu_executor.py
• pypto_serving/model/deepseek/npu_runner.py
• pypto_serving/model/deepseek/offline.py
• pypto_serving/serving/engine/async_engine.py
• pypto_serving/serving/sched/scheduler.py
• pypto_serving/serving/server/serving_worker.py
• tests/test_deepseek_v4_accuracy.py
• tests/unit/model/deepseek/test_model_components.py
• tests/unit/model/deepseek/test_offline.py
• tests/unit/serving/sched/test_async_scheduler.py
• tests/unit/serving/server/test_worker_step_protocol.py