feat(deepseek): support arbitrary MTP draft depth - #120
Open
hashiqiqixian wants to merge 12 commits into
Open
Conversation
📝 WalkthroughWalkthroughDeepSeek V4 speculative decoding now accepts configurable speculative-token depth, selects depth-specific layouts, verifies drafts in fixed-width chunks, and removes previous-token plumbing from async decode requests and batches. ChangesDeepSeek V4 speculative serving
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant Client
participant DeepSeekV4ModelRunner
participant MTPVerifier
participant DecodeExecutor
Client->>DeepSeekV4ModelRunner: submit request with speculative depth K
DeepSeekV4ModelRunner->>MTPVerifier: propose draft tokens
MTPVerifier->>DecodeExecutor: verify fixed-width target chunks
DecodeExecutor-->>MTPVerifier: return logits
MTPVerifier-->>DeepSeekV4ModelRunner: commit matching prefix and target token
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 |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (3)
pypto_serving/model/deepseek/npu_runner.py (1)
2331-2340: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winPer-step INFO acceptance log on the decode hot path.
This emits one INFO line for every decode iteration; at K=9 with a busy batch it dominates the log and adds formatting/I/O to the step. The per-request summary already lands in
release_finished_requests. Consider DEBUG or a periodic/counter-gated emit.♻️ Suggested change
- logger.info( - "DeepSeekV4 MTP acceptance progress: accepted=%d proposed=%d rate=%.2f%%", + logger.debug( + "DeepSeekV4 MTP acceptance progress: accepted=%d proposed=%d rate=%.2f%%", self._mtp_accepted_tokens, self._mtp_proposed_tokens,🤖 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 2331 - 2340, Change the acceptance-progress log in the decode path to DEBUG or gate it behind a periodic/counter-based condition so it does not emit on every iteration. Preserve the existing counters and percentage calculation, and keep per-request summary reporting in release_finished_requests unchanged.tests/test_deepseek_v4.py (1)
195-206: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueConsider covering the fully-exhausted-context case.
_mtp_draft_countcan legitimately return 0, which drives thenum_drafts == 0branch in_propose_mtp_tokensand the single-column path in_verify_mtp_drafts. Neither is exercised here; a case withseq_len == max_seq_lenwould pin that behavior.🤖 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/test_deepseek_v4.py` around lines 195 - 206, Extend test_deepseek_mtp_draft_depth_is_capped_by_remaining_context or add a focused test covering a batch sequence length equal to model.runtime.max_seq_len. Assert _mtp_draft_count returns 0, exercising the fully exhausted context and the downstream zero-draft/single-column behavior without changing existing nonzero-cap coverage.pypto_serving/cli/main.py (1)
176-177: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueError message ignores the legacy alias path.
_resolve_num_speculative_tokensmaps a bare--enable-mtpto 1, so a non-DeepSeek run with only--enable-mtpfails with a message naming a flag the user never passed. Consider mentioning both flags.♻️ Suggested wording
- elif num_speculative_tokens: - raise ValueError("--num-speculative-tokens is only supported for DeepSeek V4") + elif num_speculative_tokens: + raise ValueError( + "--num-speculative-tokens/--enable-mtp is only supported for DeepSeek V4" + )🤖 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/cli/main.py` around lines 176 - 177, Update the ValueError in the _resolve_num_speculative_tokens path to mention both --num-speculative-tokens and the legacy --enable-mtp alias, so users receive an accurate message when either flag triggers the unsupported configuration.
🤖 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/serving/server/serving_worker.py`:
- Around line 249-252: Update the `_last_tokens` field comment in `__init__` to
state that only the latest token is retained for async placeholder resolution,
matching the `recent[-1:]` behavior in the token-update method.
---
Nitpick comments:
In `@pypto_serving/cli/main.py`:
- Around line 176-177: Update the ValueError in the
_resolve_num_speculative_tokens path to mention both --num-speculative-tokens
and the legacy --enable-mtp alias, so users receive an accurate message when
either flag triggers the unsupported configuration.
In `@pypto_serving/model/deepseek/npu_runner.py`:
- Around line 2331-2340: Change the acceptance-progress log in the decode path
to DEBUG or gate it behind a periodic/counter-based condition so it does not
emit on every iteration. Preserve the existing counters and percentage
calculation, and keep per-request summary reporting in release_finished_requests
unchanged.
In `@tests/test_deepseek_v4.py`:
- Around line 195-206: Extend
test_deepseek_mtp_draft_depth_is_capped_by_remaining_context or add a focused
test covering a batch sequence length equal to model.runtime.max_seq_len. Assert
_mtp_draft_count returns 0, exercising the fully exhausted context and the
downstream zero-draft/single-column behavior without changing existing
nonzero-cap coverage.
🪄 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: 3647e323-c56e-4dc3-a924-006f9126a56a
📒 Files selected for processing (11)
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.pypypto_serving/serving/server/ipc.pypypto_serving/serving/server/serving_worker.pytests/test_batching.pytests/test_deepseek_v4.py
💤 Files with no reviewable changes (1)
- pypto_serving/config/types.py
Comment on lines
+249
to
+252
| """Remember the latest sampled token for async placeholder resolution.""" | ||
| recent = self._last_tokens.get(request_id, []) | ||
| recent.extend(int(t) for t in tokens) | ||
| if len(recent) > 2: | ||
| recent = recent[-2:] | ||
| self._last_tokens[request_id] = recent | ||
| self._last_tokens[request_id] = recent[-1:] |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Stale _last_tokens comment.
The field comment in __init__ still says "up to 2 kept for MTP prev-token context"; only the latest token is retained now.
♻️ Suggested change (lines 70-73)
- # Last-sampled tokens per request (most-recent-last, up to 2 kept for MTP
- # prev-token context). Under async scheduling the engine sends
+ # Latest sampled token per request. Under async scheduling the engine sends
# PLACEHOLDER_TOKEN for a decode input it hasn't sampled yet; the worker
# substitutes from here. Entries cleared when a request is released.🤖 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/serving/server/serving_worker.py` around lines 249 - 252,
Update the `_last_tokens` field comment in `__init__` to state that only the
latest token is retained for async placeholder resolution, matching the
`recent[-1:]` behavior in the token-update method.
added 11 commits
July 29, 2026 02:39
hashiqiqixian
force-pushed
the
feat/deepseek-v4-arbitrary-mtp
branch
from
July 28, 2026 20:17
19ec067 to
4ad65da
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
--num-speculative-tokens Kconfiguration for DeepSeek V4; any positive K enables MTP and legacy--enable-mtpremains the K=1 alias.Repository responsibility
pypto-servingowns the CLI/configuration, scheduler reservations, request/worker state, target verification, cache metadata, acceptance accounting, tests, and documentation in this PR.PyPTO-lib revision
This PR pins official
hw-native-sys/pypto-libmain revision186e1d85911f37ead07b8e74fd12734f23c947f0, which contains the DeepSeek V4 MTP projection updates used by the serving path.The PR no longer depends on hw-native-sys/pypto-lib#852. That PR was closed because there was no A/B reproduction showing its decode output-return cleanup was required by arbitrary-depth serving; its simulator stall reproduced on both attempted revisions and did not establish that dependency.
Boundaries
This PR does:
This PR does not:
Validation
Current revisions:
affe1b4f10fb2e94b5d39b706cda8f77a933cd9a186e1d85911f37ead07b8e74fd12734f23c947f0Focused CPU coverage previously passed with
97 passed, 4 warnings; the warnings were CANN installation ownership notices.Prior 8-NPU DeepSeek V4 Flash W8A8 HTTP evidence was collected before the pypto-lib dependency was removed, on serving
4ad65da806c3c42d0055dae3db91119be39b72d3with pypto-lib864b05396ec72e2bb42f5f77032a17f83e1794f8:--enable-mtp: HTTP 200, 4 completion tokens.These prior hardware results demonstrate the serving implementation but are not claimed as exact-revision validation for the current pypto-lib pin. A current-hash NPU A/B can be run separately if required.