Optimize DeepSeek V4 startup with reusable programs and prepacked weights - #124
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:
📝 WalkthroughWalkthroughAdds shared persistent kernel caching, DeepSeek V4 runtime integration, validated prepacked weight sidecars, a prepacking CLI, packaging and documentation updates, and expanded tests. ChangesDeepSeek V4 serving enhancements
Estimated code review effort: 4 (Complex) | ~60 minutes Possibly related PRs
Sequence Diagram(s)sequenceDiagram
participant Operator
participant CLI
participant DeepSeekExecutor
participant KernelCache
participant WeightStore
participant ModelRunner
Operator->>CLI: start DeepSeek V4 with cache directory
CLI->>DeepSeekExecutor: pass kernel_cache_dir
DeepSeekExecutor->>WeightStore: load validated prepacked layer weights
DeepSeekExecutor->>KernelCache: load compiled L3 program by fingerprints
KernelCache-->>DeepSeekExecutor: cached program or miss
DeepSeekExecutor->>ModelRunner: provide compiled kernels and weights
ModelRunner->>KernelCache: store newly assembled binaries
Poem
🚥 Pre-merge checks | ✅ 3 | ❌ 2❌ Failed checks (1 warning, 1 inconclusive)
✅ Passed checks (3 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 |
puddingfjz
changed the title
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 4
🧹 Nitpick comments (5)
tests/test_kernel_cache.py (1)
117-161: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueConsider covering the fingerprint-unavailable path.
Both new tests exercise the happy path. A case asserting
compute_params_fingerprint("x", object(), scalar_args=None, platform="a2a3", block_dim=None) == UNKNOWNwould pin the contract that a missing_funcdisables reuse (whichKernelCache._usabledepends on).🤖 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_kernel_cache.py` around lines 117 - 161, Add a test covering the unavailable-fingerprint path by calling compute_params_fingerprint with an object lacking _func, no scalar arguments, platform "a2a3", and no block dimension, then assert it returns UNKNOWN. Keep the existing happy-path fingerprint tests unchanged and ensure the assertion captures the contract used by KernelCache._usable.pypto_serving/tools/prepack_deepseek_v4.py (1)
103-119: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value
mkstemppublishes the sidecar with 0600 permissions.
tempfile.mkstempcreates the file mode 0600, andos.replacepreserves it, so a sidecar built by one user is unreadable by a serving process running as another (silently falling back to full repacking). Consideros.chmod(temporary_path, 0o644 & ~current_umask)— or matching the mode of a source shard — before the replace.🤖 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/tools/prepack_deepseek_v4.py` around lines 103 - 119, Update the temporary sidecar creation flow around mkstemp and os.replace to set the temporary file’s permissions to the intended readable mode before publishing it. Account for the current umask (or match the source shard mode) so the replaced output is readable by other serving users while preserving appropriate permission restrictions.pypto_serving/model/deepseek/weight_loader.py (1)
825-832: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winTensor-name drift raises instead of falling back.
Once the fingerprint matches, a name-set mismatch can only mean
_DEEPSEEK_V4_PACKED_WEIGHT_NAMESchanged without bumpingDEEPSEEK_V4_PACKED_FORMAT. Hard-failing startup for an optional optimization is harsher than the stale/cold paths, which log and repack. Either fold the name set into the fingerprint/format so old sidecars are simply "stale", or log a warning and returnNonehere.🤖 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/weight_loader.py` around lines 825 - 832, The packed-weight loader currently raises on tensor-name drift after the fingerprint matches; change the validation in the DeepSeekV4 sidecar loading path to log a warning and return None so the caller follows the existing stale/cold repack fallback. Preserve the missing/extra diagnostics in the warning, and avoid hard-failing startup for this optional optimization.tests/test_deepseek_v4.py (1)
461-498: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winPin residency in the happy-path sidecar test.
This test relies on the real
_sample_file_page_cache_residencyreporting ≥ 95% for the just-written file. On platforms withoutmincore(or under restrictive containers) it returnsNone, the sidecar is skipped, and the test then falls into real checkpoint packing and fails on thepytest.failstub. Monkeypatch it to1.0liketest_deepseek_weight_store_skips_cold_prepacked_sidecardoes.💚 Proposed fix
+ monkeypatch.setattr(weight_loader, "_sample_file_page_cache_residency", lambda path: 1.0) monkeypatch.setattr( store, "load_packed_layer_weights", lambda *args, **kwargs: pytest.fail("valid sidecar must skip checkpoint packing"), )🤖 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 461 - 498, Update test_deepseek_weight_store_maps_valid_prepacked_sidecar to monkeypatch _sample_file_page_cache_residency to return 1.0, matching test_deepseek_weight_store_skips_cold_prepacked_sidecar, so the valid sidecar path is deterministic across platforms and skips checkpoint packing.pypto_serving/model/common/kernel_cache.py (1)
145-166: 🩺 Stability & Availability | 🔵 Trivial | 💤 Low valueConcurrent
storepublish is not fully atomic despite the docstring.
shutil.rmtree(slot)followed byos.replace(tmp, slot)leaves a window where the slot is absent, and two launches storing the same name can collide (the loser'srmtree/replaceraises and is swallowed). Practically safe since a concurrentloaddegrades to a MISS, but consider softening the docstring or guarding with a lock file if parallel launches share one cache dir.🤖 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/common/kernel_cache.py` around lines 145 - 166, Update the store method’s docstring to avoid claiming the publish is fully atomic, or add synchronization for same-name concurrent stores using a lock file around the slot replacement sequence. Ensure concurrent writers do not surface rmtree/replace collisions and concurrent loads retain valid cache behavior.
🤖 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/kernel_cache.py`:
- Around line 47-76: The DeepSeek cache fingerprint must include resolved layout
and deployment dimensions so specializations cannot collide when annotations are
lazy. Update compute_params_fingerprint to incorporate decode_batch, decode_seq,
decode_tokens, and prefill_seq, and revise pypto_serving/cli/main.py lines 51-58
so its invalidation help accurately describes the dimensions covered by the
fingerprint.
In `@pypto_serving/model/deepseek/npu_executor.py`:
- Around line 369-377: Update
DeepSeekV4ModelRunner._materialize_resident_weights to clear the compiled
prepacked layer-weight reference after its tensors have been uploaded, matching
the existing embedding_weight cleanup. Release
DeepSeekV4CompiledKernels.prepacked_layer_weights (including its tensors) before
dropping _stacked_host_weights, while preserving the uploaded device weights and
released_parent_host_bytes logging.
In `@pypto_serving/model/deepseek/weight_loader.py`:
- Around line 248-299: Move the os.open call inside the guarded logic of
_sample_file_page_cache_residency so open failures are caught and return None,
preserving fallback behavior. Track the file descriptor as optional and close it
only when successfully opened. Ensure the ctypes anchor is released before
mapping.close(), and include BufferError in the handled probe failures so
cleanup issues also degrade to None rather than escaping.
In `@pypto_serving/tools/prepack_deepseek_v4.py`:
- Around line 82-95: Move the packed_stacked_layer_weights_fingerprint call
before load_stacked_layer_weights in the prepacking flow, while keeping the same
fingerprint arguments. Ensure the fingerprint is captured before the lengthy
packing operation and then used unchanged for subsequent sidecar validation.
---
Nitpick comments:
In `@pypto_serving/model/common/kernel_cache.py`:
- Around line 145-166: Update the store method’s docstring to avoid claiming the
publish is fully atomic, or add synchronization for same-name concurrent stores
using a lock file around the slot replacement sequence. Ensure concurrent
writers do not surface rmtree/replace collisions and concurrent loads retain
valid cache behavior.
In `@pypto_serving/model/deepseek/weight_loader.py`:
- Around line 825-832: The packed-weight loader currently raises on tensor-name
drift after the fingerprint matches; change the validation in the DeepSeekV4
sidecar loading path to log a warning and return None so the caller follows the
existing stale/cold repack fallback. Preserve the missing/extra diagnostics in
the warning, and avoid hard-failing startup for this optional optimization.
In `@pypto_serving/tools/prepack_deepseek_v4.py`:
- Around line 103-119: Update the temporary sidecar creation flow around mkstemp
and os.replace to set the temporary file’s permissions to the intended readable
mode before publishing it. Account for the current umask (or match the source
shard mode) so the replaced output is readable by other serving users while
preserving appropriate permission restrictions.
In `@tests/test_deepseek_v4.py`:
- Around line 461-498: Update
test_deepseek_weight_store_maps_valid_prepacked_sidecar to monkeypatch
_sample_file_page_cache_residency to return 1.0, matching
test_deepseek_weight_store_skips_cold_prepacked_sidecar, so the valid sidecar
path is deterministic across platforms and skips checkpoint packing.
In `@tests/test_kernel_cache.py`:
- Around line 117-161: Add a test covering the unavailable-fingerprint path by
calling compute_params_fingerprint with an object lacking _func, no scalar
arguments, platform "a2a3", and no block dimension, then assert it returns
UNKNOWN. Keep the existing happy-path fingerprint tests unchanged and ensure the
assertion captures the contract used by KernelCache._usable.
🪄 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: 3dd17a2a-5527-476f-8e56-ea2d942ba2ce
📒 Files selected for processing (12)
docs/dev/model/deepseek-v4.mdpyproject.tomlpypto_serving/cli/main.pypypto_serving/model/common/kernel_cache.pypypto_serving/model/deepseek/kernel_cache.pypypto_serving/model/deepseek/npu_executor.pypypto_serving/model/deepseek/npu_runner.pypypto_serving/model/deepseek/weight_loader.pypypto_serving/model/qwen/kernel_cache.pypypto_serving/tools/prepack_deepseek_v4.pytests/test_deepseek_v4.pytests/test_kernel_cache.py
Compile DeepSeek programs from signatures, reuse assembled kernel artifacts, and load validated hot prepacked weights without rebuilding the stacked layout. Harden cache fingerprints, sidecar fallback, permissions, and Host mapping cleanup.
d0cef65 to
6527c66
Compare
1 participant
Summary
Reduce DeepSeek V4 model preparation time by combining three complementary startup optimizations:
In the steady-state configuration with a warm kernel cache and hot weight sidecar, model registration decreased from 506.758 seconds to an average of 90.595 seconds: an 82.1% reduction, or 5.59x speedup.
Changes
Signature-based compilation
Persistent callable cache
--kernel-cache-dirserving argument.Prepacked layer weights
mincore().The default path remains unchanged when no kernel cache or usable sidecar is available.
Performance
Environment:
e57e14f478ddad5a9922afdb7e7d4ccHot-run component times:
The cache-hit/cold-sidecar launch ran under unusually heavy Host load: layer packing increased from the control's 90.642 seconds to 223.768 seconds. It validates the cache-hit path but is not used as a clean absolute steady-state comparison.
Sidecar behavior
The sidecar contains 322.818 GiB of prepacked weights. A cold sidecar is deliberately rejected because synchronously reading that amount of data can be slower than the original loader. The measured first full read took 577.710 seconds under concurrent machine load.
A hot mmap-backed sidecar still increases main upload from approximately 25 seconds to 50–59 seconds. Page-cache residency avoids disk I/O, but each serving process must still establish page-table mappings and pin or map approximately 84.6 million pages during the first device upload. Avoiding approximately 91 seconds of layer packing nevertheless provides a substantial net improvement.
Validation
Hardware-test limitation
The hardware launches completed program preparation, worker creation, and resident main/MTP weight upload. They then failed while sizing the KV cache because residual device allocations left only approximately 6.2 GB free, which was insufficient under
npu_memory_utilization=0.90.The reported
PyptoExecutor.register_modeltrace spans are complete through the upload and capacity-check boundary, but these measurements do not replace a clean end-to-end health, prefill, and decode validation. A follow-up run on devices without residual HBM allocations is still required.