Lower DeepSeek decode inputs to device - #129
Merged
superxf merged 5 commits intoJul 31, 2026
Merged
Conversation
📝 WalkthroughWalkthroughDeepSeek V4 decode and MTP execution now perform embedding lookup and sampling on device, use grouped cache metadata and persistent device buffers, and remove host hidden-state staging. Executor capability checks, kernel contracts, serving dispatch, resource cleanup, and tests were updated accordingly. ChangesDeepSeek device-side decode
Estimated code review effort: 4 (Complex) | ~45 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 |
high-cloud
force-pushed
the
agent/deepseek-device-metadata-serving
branch
from
July 30, 2026 07:17
eff475b to
abe01e9
Compare
high-cloud
changed the title
high-cloud
marked this pull request as ready for review
July 31, 2026 00:58
![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
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
pypto_serving/model/deepseek/npu_runner.py (1)
1965-1975: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
x_hcparameter annotation contradicts the actual argument. Both callers passx_hc=Noneand the dataclass field istorch.Tensor | None.🔧 Proposed fix
- x_hc: torch.Tensor, + x_hc: torch.Tensor | None,🤖 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 1965 - 1975, Update the x_hc parameter annotation in _prepare_decode_inputs to allow None, matching both callers and the DeepSeekV4PreparedDecodeInputs dataclass field. Preserve the existing tensor type for non-null values.
🧹 Nitpick comments (1)
pypto_serving/model/deepseek/npu_runner.py (1)
1630-1646: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valuePrefer an explicit
Nonecheck overgetattron a possibly-Nonestate.stateis a dataclass that always definestail_rank/tail_slot_id; thegetattrdefault only exists to toleratestate is None, which the next block already handles explicitly.♻️ Proposed refactor
state = self._mtp_request_states.pop(request_id, None) - tail_rank = getattr(state, "tail_rank", None) - tail_slot_id = getattr(state, "tail_slot_id", None) - if tail_rank is not None and tail_slot_id is not None: - self._mtp_free_tail_slots[tail_rank].append(tail_slot_id) - if state is not None and state.proposed_tokens: + if state is None: + continue + if state.tail_rank is not None and state.tail_slot_id is not None: + self._mtp_free_tail_slots[state.tail_rank].append(state.tail_slot_id) + if state.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 1630 - 1646, Update release_finished_requests to access state.tail_rank and state.tail_slot_id only after explicitly checking that state is not None, while preserving the existing slot-release and acceptance-logging behavior. Remove the getattr calls and keep the None guard aligned with the state dataclass’s defined fields.
🤖 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 3932-3937: Replace the ineffective hasattr(self,
"_global_weights") guard in the materialization block with a None check on the
initialized _global_weights value, matching the _inherited_host_weights pattern.
Only call load_packed_global_weights(), _materialize_embedding_device_weight(),
_materialize_main_pre_hc_device(), and the optional MTP materialization when
global weights are present.
---
Outside diff comments:
In `@pypto_serving/model/deepseek/npu_runner.py`:
- Around line 1965-1975: Update the x_hc parameter annotation in
_prepare_decode_inputs to allow None, matching both callers and the
DeepSeekV4PreparedDecodeInputs dataclass field. Preserve the existing tensor
type for non-null values.
---
Nitpick comments:
In `@pypto_serving/model/deepseek/npu_runner.py`:
- Around line 1630-1646: Update release_finished_requests to access
state.tail_rank and state.tail_slot_id only after explicitly checking that state
is not None, while preserving the existing slot-release and acceptance-logging
behavior. Remove the getattr calls and keep the None guard aligned with the
state dataclass’s defined fields.
🪄 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: 3c162216-82fc-4fbc-96f3-b2e1cf6fd8de
📒 Files selected for processing (8)
.gitignorepypto-libpypto_serving/model/common/executor/executor.pypypto_serving/model/deepseek/npu_executor.pypypto_serving/model/deepseek/npu_runner.pypypto_serving/serving/engine/engine.pypypto_serving/serving/server/serving_worker.pytests/test_deepseek_v4.py
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
2 participants
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.
PyPTO dependency
Includes merged hw-native-sys/pypto-lib#867 through upstream
pypto-librevision
95206fc.What changed
Why
DeepSeek decode still performed host embedding lookup, FP32
x_hcpacking, and MTP hidden-window construction after metadata lowering. Those transient hidden tensors added host latency and host-to-device traffic on every step.The companion PyPTO change performs embedding lookup and hidden packing inside the existing main and MTP decode graphs. No second L2 dispatch is introduced.
The checkpoint has
tie_word_embeddings=false, so the embedding table cannot reuse the LM-head weight and is uploaded independently.Impact
x_hcpackingProfile
Warm means from the exact 10-token MTP run:
2.116 ms0.035 ms0.207 ms0.159 ms38.715 ms9.161 ms51.414 msThe following former host spans are absent:
DeepSeekV4ModelRunner.decode.prepare_embeddingsDeepSeekV4ModelRunner.decode.pack_x_hcDeepSeekV4ModelRunner.mtp.lookup_embeddingsValidation
86 passed, 2 deselectedin the focused DeepSeek/batching suitegit diff --checktask_20260730_050705_331978223501task_20260730_050014_300629113280,task_20260730_050224_310317031531task_20260730_051813_38056478242a leading global information and communications technology (ICT)