Skip to content

feat(deepseek): support MTP across chunked prefill - #132

Open
vegetabledoww wants to merge 1 commit into
hw-native-sys:mainfrom
vegetabledoww:fix/deepseek-chunked-prefill-main-sync
Open

feat(deepseek): support MTP across chunked prefill#132
vegetabledoww wants to merge 1 commit into
hw-native-sys:mainfrom
vegetabledoww:fix/deepseek-chunked-prefill-main-sync

Conversation

@vegetabledoww

@vegetabledoww vegetabledoww commented Aug 1, 2026

Copy link
Copy Markdown

Summary

This PR enables DeepSeek V4 Flash W8A8 to use MTP together with chunked prefill
for prompts longer than the model's fixed 128-token prefill dispatch size.

Root cause

The scheduler had no model-specific per-request prefill limit or capability
flag for speculative decoding. DeepSeek's MTP runner also retained a whole
single-dispatch snapshot, so a long prompt could not safely carry MTP state
across chunk boundaries.

Changes

  • Add model prefill-limit and speculative-chunk compatibility fields to
    RuntimeConfig and SchedulerConfig.
  • Configure DeepSeek with a 128-token per-request prefill limit and explicitly
    advertise that MTP can span chunked prefill.
  • Split long requests at the model limit while preserving scheduler token
    budgets; reject invalid non-chunked/speculative combinations with actionable
    errors.
  • Carry one pending main-model row per request across chunks, stage padded
    per-owner MTP rows, and select the final row when the first sampled token is
    available.
  • Extend the DeepSeek kernel argument ABI, keep MoE epochs monotonic across
    prefill dispatches, and improve startup diagnostics for effective chunk size
    and MTP support.
  • Add scheduler, kernel-contract, MTP state, boundary, and cleanup regression
    coverage; update the DeepSeek development notes.

Validation

Unit tests

  • pytest -q tests/test_deepseek_v4.py: 43 passed
  • pytest -q tests/test_batching.py -k scheduler: 11 passed
  • git diff --check: passed

The combined tests/test_deepseek_v4.py tests/test_batching.py run produced
107 passing tests and 2 skips; its only failure was the unrelated existing
Qwen distributed-config compatibility check because the installed runtime still
exposes RunConfig.block_dim.

Hardware boundary check

On even devices 0,2,4,6,8,10,12,14, task
task_20260801_011207_253753715889 completed exact prompt lengths
127, 128, 129, 255, 256, 257 with MTP enabled, 100% MTP acceptance, and no
server errors.

TTFT and throughput comparison

Task task_20260801_020139_38651573533 used an exact 256-token prompt, 16
generated tokens, streaming responses, prefix caching disabled, and identical
chunk-prefill settings for the MTP and non-MTP runs:

Configuration Concurrency TTFT p50 E2E p50 Output throughput
MTP + chunk-prefill 1 430 ms 846 ms 19.0 tok/s
No MTP + chunk-prefill 1 362 ms 989 ms 16.1 tok/s
MTP + chunk-prefill 8 1,149 ms 1,955 ms 65.4 tok/s
No MTP + chunk-prefill 8 1,428 ms 2,091 ms 61.0 tok/s

MTP adds about 19% to single-request TTFT but improves single-request output
throughput by 18.4%. At concurrency 8 it reduces TTFT p50 by 19.5% and raises
aggregate output throughput by 7.1%.

The shared machine had an unrelated long-running service using devices 8-15,
so these are shared-environment measurements. The test task exited with code 0
and released all eight device locks. The companion
pypto-lib#883 must be
merged or made available first because this serving change consumes its updated
MTP prefill ABI.

Related issue

Fixes #123.

@coderabbitai

coderabbitai Bot commented Aug 1, 2026

Copy link
Copy Markdown

Review Change Stack

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: 5d9b27ee-164a-4541-82f0-251d3757c80e

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 now enforces a 128-token model prefill limit, supports automatic prompt chunking, and carries MTP state across chunks. Runtime configuration and scheduler validation expose these constraints. Runner logic updates kernel arguments, padding positions, MoE epochs, and MTP dispatch behavior.

Changes

DeepSeek V4 chunked prefill

Layer / File(s) Summary
Runtime capability and scheduler wiring
pypto_serving/config/types.py, pypto_serving/cli/main.py, pypto_serving/serving/engine/async_engine.py, docs/dev/model/deepseek-v4.md, tests/test_deepseek_v4.py
Runtime configuration records the model prefill limit and speculative chunked-prefill support. Scheduler construction and startup output expose these settings. Documentation and CLI tests describe the 128-token behavior.
Prefill admission and token limits
pypto_serving/serving/sched/scheduler.py, tests/test_batching.py
The scheduler validates prefill limits, rejects unsupported requests, caps scheduled tokens, and defers unsupported multi-step speculative prefills. Tests cover boundaries, residual chunks, disabled chunking, and residual budget handling.
Chunked MTP prefill execution
pypto_serving/model/deepseek/npu_executor.py, pypto_serving/model/deepseek/npu_runner.py, tests/test_deepseek_v4.py
MTP prefill streams pending rows across chunks. Kernel arguments include per-owner counts and MoE epochs. Padding positions, worker resets, runtime configuration, and MTP staging receive regression coverage.

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

Possibly related PRs

Poem

A rabbit hops through chunks in line,
One hundred twenty-eight at a time.
MTP rows wait, then softly flow,
While padded tails stay safe below.
MoE epochs rise in tune—
Prefill dances past the moon.

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name Status Explanation Resolution
Docstring Coverage ⚠️ Warning Docstring coverage is 21.21% 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
Title check ✅ Passed The title clearly describes the main change: enabling DeepSeek MTP across chunked prefill.
Description check ✅ Passed The description directly explains the DeepSeek chunked-prefill, MTP, scheduler, ABI, testing, and validation changes.
Linked Issues check ✅ Passed The changes address issue #123 by enforcing the 128-token limit, preserving MTP state, fixing padding positions, and adding regression coverage.
Out of Scope Changes check ✅ Passed The documented changes are focused on DeepSeek chunked prefill, MTP compatibility, scheduler limits, ABI handling, diagnostics, and related tests.

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

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

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 2

🤖 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/cli/main.py`:
- Around line 281-282: Update _validate_model_topology to reject the
incompatible --no-enable-chunked-prefill configuration for DeepSeek V4 when the
configured maximum prompt length can exceed DEEPSEEK_V4_PREFILL_SEQ (the
128-token DeepSeekV4CacheLayout().prefill_seq limit), so validation fails at
startup instead of in Scheduler.add_request. Remove the duplicate
supports_chunked_prefill_with_speculation = True assignment from the DeepSeek V4
block while preserving its earlier assignment.

In `@pypto_serving/model/deepseek/npu_runner.py`:
- Line 3896: Update the scope-statistics directory resolution near out_dir so it
no longer defaults to the developer-specific absolute path; use a neutral
temporary-directory default or require PYPTO_DSV4_SCOPE_STATS_DIR only when
PYPTO_DSV4_SCOPE_STATS enables statistics, while preserving the existing
RunConfig behavior otherwise.
🪄 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 Plus

Run ID: a28ac5dd-00a4-4c47-9326-850ef8ff0fae

📥 Commits

Reviewing files that changed from the base of the PR and between ed55def and d0474c6.

📒 Files selected for processing (9)
  • docs/dev/model/deepseek-v4.md
  • pypto_serving/cli/main.py
  • pypto_serving/config/types.py
  • pypto_serving/model/deepseek/npu_executor.py
  • pypto_serving/model/deepseek/npu_runner.py
  • pypto_serving/serving/engine/async_engine.py
  • pypto_serving/serving/sched/scheduler.py
  • tests/test_batching.py
  • tests/test_deepseek_v4.py

Comment thread pypto_serving/cli/main.py Outdated
Comment thread pypto_serving/model/deepseek/npu_runner.py Outdated
Teach the DeepSeek V4 serving scheduler and runner to split prompts at the model's 128-token prefill limit while carrying MTP state across chunk boundaries. Add explicit runtime capability metadata, padded per-owner MTP prefill staging, monotonic MoE epochs, kernel ABI arguments, and regression coverage for scheduler and request-state behavior.

Hardware validation on even devices 0,2,4,6,8,10,12,14 (task_20260801_011207_253753715889) passed prompt boundaries 127/128/129/255/256/257 with MTP enabled and 100% acceptance. Performance task_20260801_020139_38651573533 used 256 prompt tokens and 16 generated tokens: MTP+chunk reached 19.0 tok/s at concurrency 1 and 65.4 tok/s at concurrency 8, versus 16.1 and 61.0 tok/s without MTP; TTFT p50 was 430/1149 ms with MTP versus 362/1428 ms without MTP. The shared machine had an unrelated long-running service on devices 8-15, so the measurements are reported as shared-environment results.

Checks: pytest -q tests/test_deepseek_v4.py (43 passed); pytest -q tests/test_batching.py -k scheduler (11 passed); git diff --check.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[Bug] DeepSeek V4 chunked prefill does not honor the fixed 128-token kernel limit

1 participant