support kvcache dynamic shape - #115
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 cache handling now derives grouped physical capacities at runtime. Runner allocation, device cache clearing, executor tensor contracts, and serving metadata synchronization use these capacities, with tests covering sizing, allocation, retry, lifecycle, and shape behavior. ChangesDeepSeek runtime cache sizing
Estimated code review effort: 5 (Critical) | ~90 minutes Sequence Diagram(s)sequenceDiagram
participant ReplicaEngineCore
participant DeepSeekV4ModelRunner
participant KvCacheManager
participant DeepSeekV4NPUExecutor
ReplicaEngineCore->>DeepSeekV4ModelRunner: receive reported cache capacity
DeepSeekV4ModelRunner->>KvCacheManager: provide grouped physical capacity
DeepSeekV4NPUExecutor->>DeepSeekV4ModelRunner: derive runtime cache block shapes
DeepSeekV4NPUExecutor-->>DeepSeekV4ModelRunner: validate and use cache-shaped tensors
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: 3
🧹 Nitpick comments (6)
tests/test_batching.py (1)
351-373: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd a negative case for the divisibility rule.
The new validation rejects a
primary_num_blocksthat is not a multiple of the first group'smax_blocks_per_seq; that branch is currently untested, and it is exactly the path a mis-sized runner report would hit at startup.💚 Suggested additional test
def test_grouped_cache_rejects_unaligned_primary_pool(): manager = KvCacheManager(block_size=1, enable_prefix_cache=False) specs = ( KVCacheGroupSpec( name="primary", layer_indices=(0,), spec=KVCacheSpec(block_size=1, page_size_bytes=4), max_blocks_per_seq=3, ), ) with pytest.raises(ValueError, match="positive multiple"): manager.init_groups(specs, max_batch_size=8, primary_num_blocks=7)🤖 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_batching.py` around lines 351 - 373, Add a negative test alongside test_grouped_cache_capacity_scales_from_device_reported_primary_pool that constructs a primary group with max_blocks_per_seq=3 and calls init_groups with primary_num_blocks=7. Assert that init_groups raises ValueError and matches the “positive multiple” validation message.pypto_serving/model/deepseek/npu_runner.py (4)
3846-3851: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueRedundant bound check in the MTP branch.
idswere already validated againstblocks_per_layera few lines above, and this branch reuses the same bound. The MTP pool has exactly_physical_cache_num_blocks("ori")pages, so the check can be dropped.🤖 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 3846 - 3851, Remove the redundant block-ID range validation from the MTP branch near the page size calculation, including its ValueError path. Keep the earlier validation against blocks_per_layer and the existing page allocation logic unchanged.
182-186: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueClarify
decode_batchvsDEEPSEEK_V4_DECODE_BATCHin the ring divisor.
decode_batch(the caller-provided, layout-specific value) gates the minimum pool, butblocks_per_requestdivides by the module constantDEEPSEEK_V4_DECODE_BATCH. With the MTP layout (decode_batch=4) the two disagree, so the per-request ring is derived from a batch width the runtime is not using. If that is intentional (fixed kernel table depth), a one-line comment would prevent a future "fix" from changing capacity math.🤖 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 182 - 186, Clarify the intentional use of DEEPSEEK_V4_DECODE_BATCH as the divisor in blocks_per_request, noting that decode_batch is caller-provided and layout-specific while the module constant represents the fixed kernel table depth. Add a concise comment directly above the calculation without changing the existing capacity math.
1444-1474: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueOptional: fail fast on an exhausted budget and drop the defensive
getattr.
RuntimeConfigalways definesnpu_memory_utilization, sogetattr(..., 0.90)only hides typos. More usefully, whenpeak_non_kvalready exceedstotal * utilizationthe budget goes negative andcapacity_slotssilently clamps to 1, so the failure surfaces later as a generic "allocation failed at the one-slot minimum". A direct error naming the utilization shortfall would be much easier to diagnose in the field.🤖 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 1444 - 1474, Update the capacity calculation around npu_memory_utilization to access runtime.npu_memory_utilization directly instead of using getattr, and fail immediately when any device’s peak_non_kv usage exceeds its configured total-memory budget. Raise a clear error identifying the affected device and utilization shortfall before capacity_slots applies the one-slot minimum.
3822-3840: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winConsider coalescing contiguous pages into a single
copy_to.The inner loop issues one host→device transfer per (layer, block). For the
origroup that is 43 transfers per newly assigned block per rank, and this runs on the decode path via_seed_decode_work_cache_from_group_idsfor every step. Sortingidsand merging consecutive block runs (pages within a layer are contiguous) would collapse most of these into a single larger copy, with a correspondingly larger zero page. Also hoistpage_numel/page_nbytesout of the layer loop — they are invariant per tensor.🤖 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 3822 - 3840, Optimize the zero-page initialization loop in the tensor-group copy flow by hoisting the invariant page-size calculations before iterating layers, sorting ids, and coalescing consecutive block ids into contiguous runs. Issue one copy_to call per run using a sufficiently large zero buffer, while preserving the existing layer offsets, worker_id, bounds validation, and zeroing behavior in _seed_decode_work_cache_from_group_ids.pypto_serving/serving/engine/async_engine.py (1)
255-255: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valuePrefer a public initializer over reaching into
_init_blocks.The engine calls a private
KvCacheManagermethod across module boundaries. A thin public method (e.g.init_single_pool(num_blocks, block_size)) would keep the manager's API surface honest and match the grouped path'sinit_groups.🤖 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/engine/async_engine.py` at line 255, Replace the cross-module call to the private KvCacheManager._init_blocks in the async engine with a public initializer, such as init_single_pool(num_blocks, block_size), and implement or reuse that method on KvCacheManager to delegate to the existing single-pool initialization behavior. Keep the grouped initialization path using init_groups.
🤖 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_executor.py`:
- Around line 585-591: The docstring for the shape-producing method should stop
describing the returned dimensions as “worst-case.” Reword it to identify them
as minimum-sized or lower-bound JIT dummy shapes, while preserving the
explanation that runtime HBM-backed allocations determine actual dimensions and
generated kernels must derive strides from runtime tensor shapes.
In `@pypto_serving/model/deepseek/npu_runner.py`:
- Around line 4244-4258: Update _ensure_cache_zero_page to explicitly validate
that self._cache_group_specs is non-empty before computing max_page_bytes. Raise
a clear, named error describing that cache group specs must be resolved before
creating the zero page, while preserving the existing calculation for valid
specs.
- Around line 1416-1430: Update _fallback_cache_group_num_blocks to derive one
shared capacity_slots value from the layout’s physical limits and decode_batch,
rounded down to a whole number of group_specs[0].max_blocks_per_seq slots.
Return each cache group’s capacity as capacity_slots multiplied by its
corresponding max_blocks_per_seq, ensuring all fallback values are slot-aligned
and use consistent engine-side scaling.
---
Nitpick comments:
In `@pypto_serving/model/deepseek/npu_runner.py`:
- Around line 3846-3851: Remove the redundant block-ID range validation from the
MTP branch near the page size calculation, including its ValueError path. Keep
the earlier validation against blocks_per_layer and the existing page allocation
logic unchanged.
- Around line 182-186: Clarify the intentional use of DEEPSEEK_V4_DECODE_BATCH
as the divisor in blocks_per_request, noting that decode_batch is
caller-provided and layout-specific while the module constant represents the
fixed kernel table depth. Add a concise comment directly above the calculation
without changing the existing capacity math.
- Around line 1444-1474: Update the capacity calculation around
npu_memory_utilization to access runtime.npu_memory_utilization directly instead
of using getattr, and fail immediately when any device’s peak_non_kv usage
exceeds its configured total-memory budget. Raise a clear error identifying the
affected device and utilization shortfall before capacity_slots applies the
one-slot minimum.
- Around line 3822-3840: Optimize the zero-page initialization loop in the
tensor-group copy flow by hoisting the invariant page-size calculations before
iterating layers, sorting ids, and coalescing consecutive block ids into
contiguous runs. Issue one copy_to call per run using a sufficiently large zero
buffer, while preserving the existing layer offsets, worker_id, bounds
validation, and zeroing behavior in _seed_decode_work_cache_from_group_ids.
In `@pypto_serving/serving/engine/async_engine.py`:
- Line 255: Replace the cross-module call to the private
KvCacheManager._init_blocks in the async engine with a public initializer, such
as init_single_pool(num_blocks, block_size), and implement or reuse that method
on KvCacheManager to delegate to the existing single-pool initialization
behavior. Keep the grouped initialization path using init_groups.
In `@tests/test_batching.py`:
- Around line 351-373: Add a negative test alongside
test_grouped_cache_capacity_scales_from_device_reported_primary_pool that
constructs a primary group with max_blocks_per_seq=3 and calls init_groups with
primary_num_blocks=7. Assert that init_groups raises ValueError and matches the
“positive multiple” validation message.
🪄 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: f9864683-4eba-422e-8a2f-c11390662d45
📒 Files selected for processing (7)
pypto-libpypto_serving/model/deepseek/npu_executor.pypypto_serving/model/deepseek/npu_runner.pypypto_serving/serving/engine/async_engine.pypypto_serving/serving/memory/kv_cache.pytests/test_batching.pytests/test_deepseek_v4.py
| """Return worst-case runtime pool shapes used for JIT dummy arguments. | ||
|
|
||
| Actual worker allocations are HBM-sized after compilation. Like Qwen's | ||
| paged kernels, DeepSeek kernels must derive their layer/page stride from | ||
| the runtime tensor shape rather than baking these dummy dimensions into | ||
| generated code. | ||
| """ |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Docstring says "worst-case" but this is a minimum-sized shape.
local_capacity_slots is ceil(max_batch_size / ranks), while the runtime pool is HBM-derived and normally far larger. These dummy dimensions are a lower bound, not a worst case; the surrounding note about deriving strides from the runtime shape is the real invariant. Reword so nobody treats this value as an upper bound on the pool.
🤖 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_executor.py` around lines 585 - 591, The
docstring for the shape-producing method should stop describing the returned
dimensions as “worst-case.” Reword it to identify them as minimum-sized or
lower-bound JIT dummy shapes, while preserving the explanation that runtime
HBM-backed allocations determine actual dimensions and generated kernels must
derive strides from runtime tensor shapes.
| def _fallback_cache_group_num_blocks(self) -> dict[str, int]: | ||
| """Return the old fixed capacities for shape-only/non-runtime callers.""" | ||
| layout = self._compiled.layout | ||
| physical_limits = { | ||
| "ori": layout.decode_ori_max_blocks, | ||
| "cmp": layout.cmp_max_blocks, | ||
| "idx": layout.idx_max_blocks, | ||
| "hca_state": layout.hca_state_max_blocks, | ||
| "csa_state": layout.csa_state_max_blocks, | ||
| "csa_inner_state": layout.csa_inner_state_max_blocks, | ||
| } | ||
| return { | ||
| name: max(1, physical_limits[name] - layout.decode_batch) | ||
| for name in DEEPSEEK_V4_CACHE_GROUP_NAMES | ||
| } |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
Fallback capacities are not slot-aligned, which breaks the engine-side contract.
init_kv_cache returns _cache_group_num_blocks["ori"] and ReplicaEngineCore.start feeds it to KvCacheManager.init_groups(primary_num_blocks=...), which rejects any value that is not a multiple of group_specs[0].max_blocks_per_seq. The fallback computes physical_limit - decode_batch per group, e.g. 128 - 8 = 120 against max_blocks_per_seq = 128 // 8 = 16 → not a multiple, so startup raises. Note this path is reachable outside unit tests: the executor always sets runtime_model, but l3_callables() is empty whenever compile_kernels=False, so a real worker can land here. Additionally the per-group fallback values are not derived from one shared capacity_slots, so the engine's slots * max_blocks_per_seq scaling would not match the runner's pools even when divisibility happens to hold.
🛠️ Derive the fallback from a whole number of capacity slots
def _fallback_cache_group_num_blocks(self) -> dict[str, int]:
"""Return the old fixed capacities for shape-only/non-runtime callers."""
layout = self._compiled.layout
physical_limits = {
"ori": layout.decode_ori_max_blocks,
"cmp": layout.cmp_max_blocks,
"idx": layout.idx_max_blocks,
"hca_state": layout.hca_state_max_blocks,
"csa_state": layout.csa_state_max_blocks,
"csa_inner_state": layout.csa_inner_state_max_blocks,
}
- return {
- name: max(1, physical_limits[name] - layout.decode_batch)
- for name in DEEPSEEK_V4_CACHE_GROUP_NAMES
- }
+ specs = {spec.name: spec for spec in self._cache_group_specs}
+ slots = max(
+ 1,
+ min(
+ max(1, physical_limits[name] - layout.decode_batch)
+ // specs[name].max_blocks_per_seq
+ for name in DEEPSEEK_V4_CACHE_GROUP_NAMES
+ ),
+ )
+ return deepseek_v4_cache_blocks_for_slots(self._cache_group_specs, slots)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| def _fallback_cache_group_num_blocks(self) -> dict[str, int]: | |
| """Return the old fixed capacities for shape-only/non-runtime callers.""" | |
| layout = self._compiled.layout | |
| physical_limits = { | |
| "ori": layout.decode_ori_max_blocks, | |
| "cmp": layout.cmp_max_blocks, | |
| "idx": layout.idx_max_blocks, | |
| "hca_state": layout.hca_state_max_blocks, | |
| "csa_state": layout.csa_state_max_blocks, | |
| "csa_inner_state": layout.csa_inner_state_max_blocks, | |
| } | |
| return { | |
| name: max(1, physical_limits[name] - layout.decode_batch) | |
| for name in DEEPSEEK_V4_CACHE_GROUP_NAMES | |
| } | |
| def _fallback_cache_group_num_blocks(self) -> dict[str, int]: | |
| """Return the old fixed capacities for shape-only/non-runtime callers.""" | |
| layout = self._compiled.layout | |
| physical_limits = { | |
| "ori": layout.decode_ori_max_blocks, | |
| "cmp": layout.cmp_max_blocks, | |
| "idx": layout.idx_max_blocks, | |
| "hca_state": layout.hca_state_max_blocks, | |
| "csa_state": layout.csa_state_max_blocks, | |
| "csa_inner_state": layout.csa_inner_state_max_blocks, | |
| } | |
| specs = {spec.name: spec for spec in self._cache_group_specs} | |
| slots = max( | |
| 1, | |
| min( | |
| max(1, physical_limits[name] - layout.decode_batch) | |
| // specs[name].max_blocks_per_seq | |
| for name in DEEPSEEK_V4_CACHE_GROUP_NAMES | |
| ), | |
| ) | |
| return deepseek_v4_cache_blocks_for_slots(self._cache_group_specs, slots) |
🤖 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 1416 - 1430, Update
_fallback_cache_group_num_blocks to derive one shared capacity_slots value from
the layout’s physical limits and decode_batch, rounded down to a whole number of
group_specs[0].max_blocks_per_seq slots. Return each cache group’s capacity as
capacity_slots multiplied by its corresponding max_blocks_per_seq, ensuring all
fallback values are slot-aligned and use consistent engine-side scaling.
| def _ensure_cache_zero_page(self) -> torch.Tensor: | ||
| """Create one pre-fork zero source reused for device-page clearing.""" | ||
| if self._cache_zero_page is not None: | ||
| return self._cache_zero_page | ||
| self._ensure_shared_host_allocation_before_worker("cache zero page") | ||
| max_page_bytes = max( | ||
| spec.spec.page_size_bytes // max(len(spec.layer_indices), 1) | ||
| for spec in self._cache_group_specs | ||
| ) | ||
| for tensor in ( | ||
| cache.kv_cache, | ||
| cache.cmp_kv, | ||
| cache.idx_kv_cache, | ||
| cache.idx_kv_scale, | ||
| cache.hca_compress_state, | ||
| cache.csa_compress_state, | ||
| cache.csa_inner_compress_state, | ||
| ): | ||
| tensor.zero_() | ||
| self._decode_work_cache = cache | ||
| return cache | ||
| max_page_bytes = max( | ||
| max_page_bytes, | ||
| self._compiled.layout.block_size | ||
| * DEEPSEEK_V4_HEAD_DIM | ||
| * torch.bfloat16.itemsize, | ||
| ) |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win
Guard the empty-spec case with an explicit error.
max(...) over self._cache_group_specs raises ValueError: max() arg is an empty sequence when specs have not been resolved yet. That ordering is reachable: preflight() calls _ensure_l3_shared_buffers directly, which now stages the zero page, and _cache_group_specs is only populated by init_kv_cache. A named error makes the ordering requirement obvious.
🛡️ Proposed guard
self._ensure_shared_host_allocation_before_worker("cache zero page")
+ if not self._cache_group_specs:
+ raise RuntimeError(
+ "DeepSeekV4 cache group specs must be resolved before staging the zero page"
+ )
max_page_bytes = max(📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| def _ensure_cache_zero_page(self) -> torch.Tensor: | |
| """Create one pre-fork zero source reused for device-page clearing.""" | |
| if self._cache_zero_page is not None: | |
| return self._cache_zero_page | |
| self._ensure_shared_host_allocation_before_worker("cache zero page") | |
| max_page_bytes = max( | |
| spec.spec.page_size_bytes // max(len(spec.layer_indices), 1) | |
| for spec in self._cache_group_specs | |
| ) | |
| for tensor in ( | |
| cache.kv_cache, | |
| cache.cmp_kv, | |
| cache.idx_kv_cache, | |
| cache.idx_kv_scale, | |
| cache.hca_compress_state, | |
| cache.csa_compress_state, | |
| cache.csa_inner_compress_state, | |
| ): | |
| tensor.zero_() | |
| self._decode_work_cache = cache | |
| return cache | |
| max_page_bytes = max( | |
| max_page_bytes, | |
| self._compiled.layout.block_size | |
| * DEEPSEEK_V4_HEAD_DIM | |
| * torch.bfloat16.itemsize, | |
| ) | |
| def _ensure_cache_zero_page(self) -> torch.Tensor: | |
| """Create one pre-fork zero source reused for device-page clearing.""" | |
| if self._cache_zero_page is not None: | |
| return self._cache_zero_page | |
| self._ensure_shared_host_allocation_before_worker("cache zero page") | |
| if not self._cache_group_specs: | |
| raise RuntimeError( | |
| "DeepSeekV4 cache group specs must be resolved before staging the zero page" | |
| ) | |
| max_page_bytes = max( | |
| spec.spec.page_size_bytes // max(len(spec.layer_indices), 1) | |
| for spec in self._cache_group_specs | |
| ) | |
| max_page_bytes = max( | |
| max_page_bytes, | |
| self._compiled.layout.block_size | |
| * DEEPSEEK_V4_HEAD_DIM | |
| * torch.bfloat16.itemsize, | |
| ) |
🤖 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 4244 - 4258, Update
_ensure_cache_zero_page to explicitly validate that self._cache_group_specs is
non-empty before computing max_page_bytes. Raise a clear, named error describing
that cache group specs must be resolved before creating the zero page, while
preserving the existing calculation for valid specs.
be7d3ec to
68e288c
Compare
| if primary_num_blocks is not None: | ||
| primary_num_blocks = int(primary_num_blocks) | ||
| primary_stride = group_specs[0].max_blocks_per_seq | ||
| if primary_num_blocks <= 0 or primary_num_blocks % primary_stride: |
There was a problem hiding this comment.
_fallback_cache_group_num_blocks 里计算的大小不符合整除要求，是否要一起修改
There was a problem hiding this comment.
已修改，先会进行取整
capacity_slots = min( (physical_limit - decode_batch) // max_blocks_per_seq )
| ) | ||
| else: | ||
| self.kv_cache_manager._init_blocks(actual_num_pages, self._runtime.page_size) | ||
| self.kv_cache_manager._init_blocks(reported_num_blocks, self._runtime.page_size) |
There was a problem hiding this comment.
KvCacheManager 能否提供统一、public的初始化，自己内部区分 init_groups 或 _init_blocks
There was a problem hiding this comment.
之后在 KvCacheManager 增加统一 public 初始化接口，由 manager 根据 RuntimeConfig.kv_cache_groups 选择对应路径
2 participants
Uh oh!
There was an error while loading. Please reload this page.