opt deepseek host time

[superxf]

opt deepseek host time

[coderabbitai]

📝 Walkthrough

Walkthrough

DeepSeek V4 decode and MTP preparation now use persistent shared buffers, caller-provided packed outputs, cached static metadata, centralized MTP staging, and lifecycle state resets. Tests cover metadata reuse and active-prefix updates.

Changes

DeepSeek decode and MTP staging

Layer / File(s)	Summary
Shared decode packing pypto_serving/model/deepseek/npu_runner.py	Decode and MTP packers accept shared output tensors, validate them, and control whether all packed regions are initialized.
Reusable decode preparation pypto_serving/model/deepseek/npu_runner.py, tests/test_deepseek_v4.py	Decode inputs and metadata are written into persistent buffers, static ring-table metadata is cached by padded group block IDs, and returned tensors reference reusable storage.
MTP staging and lifecycle pypto_serving/model/deepseek/npu_runner.py, tests/test_deepseek_v4.py	MTP staging is centralized with active-prefix updates, decode dispatch avoids host-side output clearing, profiling spans are added, and initialization caches reset when buffers are recreated or closed.

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

Sequence Diagram(s)

sequenceDiagram
  participant Request as Decode request
  participant Runner as DeepSeekV4ModelRunner
  participant Builder as DeepSeekV4InputBuilder
  participant Buffers as Shared decode buffers
  participant Metadata as Metadata cache
  Request->>Runner: prepare decode inputs
  Runner->>Builder: pack x_hc into shared output
  Builder->>Buffers: write packed token rows
  Runner->>Metadata: reuse or rebuild static metadata
  Metadata-->>Runner: block-table metadata
  Runner-->>Request: prepared buffer views

Loading

Possibly related PRs

• hw-native-sys/pypto-serving#89: Related MTP staging and dispatch buffer changes.
• hw-native-sys/pypto-serving#117: Related MTP draft selection and shared-buffer flow.
• hw-native-sys/pypto-serving#120: Related MTP draft-depth and verification pipeline changes.

Poem

A rabbit hops through buffers bright,
Packing rows from morn till night.
Metadata rests, then wakes with care,
MTP leaves less work in air.
Shared tensors twitch, the kernels cheer—
“Reuse the space from batch to batch!” 🐇

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title matches the main goal of optimizing DeepSeek host time, even if it is abbreviated.
Description check	✅ Passed	The description is related to the pull request and reflects the same DeepSeek host-time optimization goal.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
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.

---

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 (3)

tests/test_deepseek_v4.py (1)

1284-1288: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Match production shape for decode_logit_row_indices.

_ensure_mtp_buffers allocates this as (ranks, DEEPSEEK_V4_MAX_LOGIT_ROWS); the stub uses (ranks, layout.decode_tokens). It happens to be large enough for local_row < decode_batch, but the divergence would hide an indexing regression against the real buffer.

♻️ Proposed stub fix

         decode_logit_row_indices=torch.empty(
             layout.ranks,
-            layout.decode_tokens,
+            DEEPSEEK_V4_MAX_LOGIT_ROWS,
             dtype=torch.int32,
         ),

🤖 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 1284 - 1288, Update the test stub’s
decode_logit_row_indices allocation to match the production shape used by
_ensure_mtp_buffers: use layout.ranks by DEEPSEEK_V4_MAX_LOGIT_ROWS instead of
layout.decode_tokens. Preserve the existing dtype and allocation behavior so
indexing tests exercise the real buffer dimensions.

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

1942-1964: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

x_hc plumbing into _prepare_decode_inputs is now redundant.

Both callers pack directly into buffers.x_hc_a and then pass the same tensor down, where _prepare_decode_inputs re-stages it via _copy_shared (Line 2037) — a no-op only because of the new data_ptr short-circuit. Dropping the x_hc parameter and the stage_x_hc span would remove the dead round-trip and make the ownership contract explicit.

Also applies to: 1979-2007

🤖 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 1942 - 1964, Remove
the redundant x_hc argument from both decode call sites and the
_prepare_decode_inputs signature. Delete the stage_x_hc profile span and
_copy_shared round-trip there, using buffers.x_hc_a directly as the
already-packed input. Keep both callers’ packing into buffers.x_hc_a and
preserve the existing decode behavior.

---

984-1044: 🩺 Stability & Availability | 🔵 Trivial | 💤 Low value

Validate local_rows before computing fill_tokens.

With fill_all=False, fill_tokens is derived from max(local_rows) at Line 1022, but the 0 <= local_row < decode_batch check only happens later in the loop at Line 1036. An out-of-range local_row makes the prefix slice silently shorter than requested and surfaces as an opaque copy_/expand shape error instead of the intended ValueError.

🛡️ Proposed validation hoist

         fallback = token_rows[0, -1]
         fill_tokens = self.layout.decode_tokens
         if not fill_all:
+            for local_row in local_rows:
+                if not 0 <= int(local_row) < self.layout.decode_batch:
+                    raise ValueError(f"decode local row {int(local_row)} is out of range")
             fill_tokens = (max(int(row) for row in local_rows) + 1) * self.layout.decode_seq

🤖 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 984 - 1044, Validate
every local row in local_rows before computing fill_tokens in the out is not
None path, reusing the existing decode local-row bounds check and preserving the
intended ValueError for invalid values. Then derive the fill_tokens prefix only
after validation, while retaining the per-entry rank validation and copy
behavior in the existing loop.

🤖 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 1942-1964: Remove the redundant x_hc argument from both decode
call sites and the _prepare_decode_inputs signature. Delete the stage_x_hc
profile span and _copy_shared round-trip there, using buffers.x_hc_a directly as
the already-packed input. Keep both callers’ packing into buffers.x_hc_a and
preserve the existing decode behavior.
- Around line 984-1044: Validate every local row in local_rows before computing
fill_tokens in the out is not None path, reusing the existing decode local-row
bounds check and preserving the intended ValueError for invalid values. Then
derive the fill_tokens prefix only after validation, while retaining the
per-entry rank validation and copy behavior in the existing loop.

In `@tests/test_deepseek_v4.py`:
- Around line 1284-1288: Update the test stub’s decode_logit_row_indices
allocation to match the production shape used by _ensure_mtp_buffers: use
layout.ranks by DEEPSEEK_V4_MAX_LOGIT_ROWS instead of layout.decode_tokens.
Preserve the existing dtype and allocation behavior so indexing tests exercise
the real buffer dimensions.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro Plus

Run ID: a136119e-4fb6-49f4-9c7a-ca2a0bde6257

📥 Commits

Reviewing files that changed from the base of the PR and between e57e14f and 6293117.

📒 Files selected for processing (2)

• pypto_serving/model/deepseek/npu_runner.py
• tests/test_deepseek_v4.py