feat(deepseek): support MTP across chunked prefill - #132
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:
📝 WalkthroughWalkthroughDeepSeek 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. ChangesDeepSeek V4 chunked prefill
Estimated code review effort: 4 (Complex) | ~60 minutes Possibly related PRs
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 |
There was a problem hiding this comment.
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
📒 Files selected for processing (9)
docs/dev/model/deepseek-v4.mdpypto_serving/cli/main.pypypto_serving/config/types.pypypto_serving/model/deepseek/npu_executor.pypypto_serving/model/deepseek/npu_runner.pypypto_serving/serving/engine/async_engine.pypypto_serving/serving/sched/scheduler.pytests/test_batching.pytests/test_deepseek_v4.py
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.
41d8f33 to
d118eac
Compare
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
RuntimeConfigandSchedulerConfig.advertise that MTP can span chunked prefill.
budgets; reject invalid non-chunked/speculative combinations with actionable
errors.
per-owner MTP rows, and select the final row when the first sampled token is
available.
prefill dispatches, and improve startup diagnostics for effective chunk size
and MTP support.
coverage; update the DeepSeek development notes.
Validation
Unit tests
pytest -q tests/test_deepseek_v4.py: 43 passedpytest -q tests/test_batching.py -k scheduler: 11 passedgit diff --check: passedThe combined
tests/test_deepseek_v4.py tests/test_batching.pyrun produced107 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, tasktask_20260801_011207_253753715889completed exact prompt lengths127, 128, 129, 255, 256, 257with MTP enabled, 100% MTP acceptance, and noserver errors.
TTFT and throughput comparison
Task
task_20260801_020139_38651573533used an exact 256-token prompt, 16generated tokens, streaming responses, prefix caching disabled, and identical
chunk-prefill settings for the MTP and non-MTP runs:
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.