deepseek/v4: align cache layout constants#674
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces centralized static paged-cache pool constants in config.py and refactors various decode and prefill attention/compressor modules to use these unified configuration variables instead of hardcoded values. Additionally, it removes the lm_head_tp call and its associated buffer allocations from both decode_fwd.py and prefill_fwd.py, returning hidden_out directly instead of logits. Feedback on these changes suggests updating an outdated comment in config.py regarding the receiver buffer floor, removing the now-unused lm_head_weight parameter and imports from the forward pass files, and adding an assertion in prefill_indexer.py to enforce configuration invariants for INDEXER_SCORE_CAP.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
|
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 Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
📝 WalkthroughWalkthroughThis PR centralizes static KV/index cache pool sizing (original, compressed, index) as config constants for decode and prefill modes, replacing hardcoded or locally derived values across many attention, compressor, and indexer kernels. Separately, decode and prefill forward drivers stop computing logits via lm_head and instead output hidden states directly. ChangesConfig-driven KV/index cache sizing
Decode/Prefill forward hidden_out output
Estimated code review effort: 4 (Complex) | ~60 minutes Possibly related PRs
Suggested labels: Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 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: 2
🧹 Nitpick comments (4)
models/deepseek/v4/config.py (2)
284-285: 🧹 Nitpick | 🔵 TrivialRecv buffer floor increase (16 → 1024) raises fixed memory footprint.
This is a 64x increase to the minimum recv-buffer size regardless of actual token/expert-derived need, applied to both decode and prefill. Confirm this floor was sized for the current distributed layout's worst case and not just a generously safe round number, since it directly scales per-rank buffer allocation.
🤖 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 `@models/deepseek/v4/config.py` around lines 284 - 285, The recv buffer minimums in the DeepSeek v4 config are being raised from a small floor to 1024 for both decode and prefill, which increases fixed per-rank memory regardless of the token-derived sizing. Review the buffer sizing logic around DECODE_RECV_MAX and PREFILL_RECV_MAX in config.py, confirm the floor is actually required for the current EP/FLASH distributed layout worst case, and if not, reduce it or tie it to a justified layout-specific bound instead of a generic round number.
265-270: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winPrefill CMP/IDX block counts bypass the batch-multiply pattern used for ORI.
PREFILL_ORI_BLOCK_NUMfollowsPREFILL_BATCH * PREFILL_ORI_MAX_BLOCKS, butPREFILL_CMP_BLOCK_NUM/PREFILL_IDX_BLOCK_NUMare set directly to the decode totals (DECODE_CMP_BLOCK_NUM/DECODE_IDX_BLOCK_NUM) rather thanPREFILL_BATCH * PREFILL_CMP_MAX_BLOCKS(which would yield 8, not 32). The comment at Lines 256-258 explains this is intentional (shared decode-batch physical pool), but the inconsistent naming/formula pattern between ORI and CMP/IDX is easy to misread as a bug during future maintenance. Consider a clearer name (e.g.,PREFILL_CMP_BLOCK_NUM = DECODE_CMP_BLOCK_NUM # shared decode-batch pool, not PREFILL_BATCH-scaled) or an inline comment at each assignment to make the divergence explicit at the point of use.🤖 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 `@models/deepseek/v4/config.py` around lines 265 - 270, The PREFILL CMP/IDX block-count assignments in config.py intentionally diverge from the ORI batch-scaled pattern, so make that explicit where the constants are defined. Update the PREFILL_CMP_BLOCK_NUM and PREFILL_IDX_BLOCK_NUM assignments near the PREFILL_ORI_MAX_BLOCKS/PREFILL_ORI_BLOCK_NUM block to clearly indicate they reuse the decode pool (matching DECODE_CMP_BLOCK_NUM and DECODE_IDX_BLOCK_NUM) rather than multiplying by PREFILL_BATCH, using either clearer naming or inline comments. Keep the existing ORI formula unchanged and ensure the intent is obvious in the same section as PREFILL_CMP_MAX_BLOCKS and PREFILL_IDX_MAX_BLOCKS.models/deepseek/v4/decode_attention_csa.py (1)
93-97: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winInconsistent migration:
ORI_BLOCK_NUMstill locally derived instead of usingDECODE_ORI_BLOCK_NUM.
CMP_BLOCK_NUMandIDX_CACHE_BLOCK_NUMare now sourced directly from the shared config constants (DECODE_CMP_BLOCK_NUM,DECODE_IDX_BLOCK_NUM), butORI_BLOCK_NUMis still locally re-derived asB * ORI_MAX_BLOCKSinstead of using the sibling constantDECODE_ORI_BLOCK_NUM(already defined in config and used this way indecode_attention_hca.py,decode_sparse_attn_hca.py, anddecode_sparse_attn_swa.py). While numerically equivalent today (assumingB == DECODE_BATCH), this leaves one path un-migrated and could silently diverge if the config-side derivation changes.♻️ Proposed fix
from config import ( DECODE_CMP_BLOCK_NUM, DECODE_IDX_BLOCK_NUM, + DECODE_ORI_BLOCK_NUM, IDX_CACHE_MAX_BLOCKS, KV_CMP_MAX_BLOCKS, KV_ORI_MAX_BLOCKS, ) ... IDX_CACHE_BLOCK_NUM = DECODE_IDX_BLOCK_NUM ORI_MAX_BLOCKS = KV_ORI_MAX_BLOCKS -ORI_BLOCK_NUM = B * ORI_MAX_BLOCKS +ORI_BLOCK_NUM = DECODE_ORI_BLOCK_NUM CMP_MAX_BLOCKS = KV_CMP_MAX_BLOCKS CMP_BLOCK_NUM = DECODE_CMP_BLOCK_NUM🤖 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 `@models/deepseek/v4/decode_attention_csa.py` around lines 93 - 97, `ORI_BLOCK_NUM` in `decode_attention_csa.py` is still derived locally from `B * ORI_MAX_BLOCKS` while the other block counts already use shared config constants. Update the `decode_attention_csa.py` constants block to source `ORI_BLOCK_NUM` from `DECODE_ORI_BLOCK_NUM`, matching the pattern used in `decode_attention_hca.py`, `decode_sparse_attn_hca.py`, and `decode_sparse_attn_swa.py`. Keep the existing `ORI_MAX_BLOCKS` alias if needed, but ensure `ORI_BLOCK_NUM` is no longer independently computed so all decode paths stay aligned with config.models/deepseek/v4/prefill_attention_hca.py (1)
75-92: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winModule-level asserts are now tautological.
SPARSE_ORI_BLOCK_NUMis directly assignedB * SPARSE_ORI_MAX_BLOCKS(line 76) andSPARSE_CMP_BLOCK_NUMis directly assignedPREFILL_CMP_BLOCK_NUM(line 79). The subsequent guards:assert SPARSE_ORI_BLOCK_NUM == B * SPARSE_ORI_MAX_BLOCKS assert SPARSE_CMP_BLOCK_NUM == PREFILL_CMP_BLOCK_NUMnow compare each value against the exact expression used to define it, so they can never fail — they no longer catch a real divergence (e.g. if
PREFILL_CMP_BLOCK_NUMin config.py were changed to something inconsistent with the geometry these sparse caches actually need). Before this refactor these constants were presumably derived independently via ceil/rounding formulas, and the asserts meaningfully cross-checked that against config; that safety net is now gone.Consider removing these two asserts (they add no value) or replacing them with a check against an independently-derived expectation (e.g. re-deriving the previous ceil/rounding formula and asserting it still matches the config value), so the guard still catches real config/geometry drift.
🔧 Suggested cleanup
assert S == COMPRESS_RATIO, "first prefill HCA bring-up targets one ratio-128 prompt chunk" assert WIN == BLOCK_SIZE, "prefill HCA currently assumes one window page per batch" -assert SPARSE_ORI_BLOCK_NUM == B * SPARSE_ORI_MAX_BLOCKS -assert SPARSE_CMP_BLOCK_NUM == PREFILL_CMP_BLOCK_NUM🤖 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 `@models/deepseek/v4/prefill_attention_hca.py` around lines 75 - 92, The module-level assertions in prefill_attention_hca.py are tautological because SPARSE_ORI_BLOCK_NUM and SPARSE_CMP_BLOCK_NUM are now assigned directly from the same expressions they assert against. Update the guard logic around SPARSE_ORI_BLOCK_NUM, SPARSE_CMP_BLOCK_NUM, and PREFILL_CMP_BLOCK_NUM to either remove the redundant asserts or replace them with independently-derived expectations that still validate the sparse-cache geometry against config drift.
🤖 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 `@models/deepseek/v4/decode_fwd.py`:
- Around line 700-701: Remove the stale lm_head_weight requirement from the
hidden-state contract in decode_fwd.py: the L3 path now only produces hidden_out
and should no longer accept or order lm_head_weight. Update the relevant
spec/typing declaration and the ordered_names logic to drop "lm_head_weight",
and delete the now-unreachable LM_HEAD_NAMES branch in the contract builder if
nothing else references it. Keep the hidden_out path intact so callers using the
serving-side LM-head flow no longer need to pass the unused tensor.
In `@models/deepseek/v4/prefill_fwd.py`:
- Around line 759-760: The hidden-state contract still exposes a stale
lm_head_weight input even though l3_prefill_fwd now returns hidden_out directly
and no longer uses lm_head_tp. Remove lm_head_weight from the relevant specs and
update ordered_names in prefill_fwd.py so the exported contract matches the new
serving-side path; also clean up the LM_HEAD_NAMES branch if it is now
unreachable.
---
Nitpick comments:
In `@models/deepseek/v4/config.py`:
- Around line 284-285: The recv buffer minimums in the DeepSeek v4 config are
being raised from a small floor to 1024 for both decode and prefill, which
increases fixed per-rank memory regardless of the token-derived sizing. Review
the buffer sizing logic around DECODE_RECV_MAX and PREFILL_RECV_MAX in
config.py, confirm the floor is actually required for the current EP/FLASH
distributed layout worst case, and if not, reduce it or tie it to a justified
layout-specific bound instead of a generic round number.
- Around line 265-270: The PREFILL CMP/IDX block-count assignments in config.py
intentionally diverge from the ORI batch-scaled pattern, so make that explicit
where the constants are defined. Update the PREFILL_CMP_BLOCK_NUM and
PREFILL_IDX_BLOCK_NUM assignments near the
PREFILL_ORI_MAX_BLOCKS/PREFILL_ORI_BLOCK_NUM block to clearly indicate they
reuse the decode pool (matching DECODE_CMP_BLOCK_NUM and DECODE_IDX_BLOCK_NUM)
rather than multiplying by PREFILL_BATCH, using either clearer naming or inline
comments. Keep the existing ORI formula unchanged and ensure the intent is
obvious in the same section as PREFILL_CMP_MAX_BLOCKS and
PREFILL_IDX_MAX_BLOCKS.
In `@models/deepseek/v4/decode_attention_csa.py`:
- Around line 93-97: `ORI_BLOCK_NUM` in `decode_attention_csa.py` is still
derived locally from `B * ORI_MAX_BLOCKS` while the other block counts already
use shared config constants. Update the `decode_attention_csa.py` constants
block to source `ORI_BLOCK_NUM` from `DECODE_ORI_BLOCK_NUM`, matching the
pattern used in `decode_attention_hca.py`, `decode_sparse_attn_hca.py`, and
`decode_sparse_attn_swa.py`. Keep the existing `ORI_MAX_BLOCKS` alias if needed,
but ensure `ORI_BLOCK_NUM` is no longer independently computed so all decode
paths stay aligned with config.
In `@models/deepseek/v4/prefill_attention_hca.py`:
- Around line 75-92: The module-level assertions in prefill_attention_hca.py are
tautological because SPARSE_ORI_BLOCK_NUM and SPARSE_CMP_BLOCK_NUM are now
assigned directly from the same expressions they assert against. Update the
guard logic around SPARSE_ORI_BLOCK_NUM, SPARSE_CMP_BLOCK_NUM, and
PREFILL_CMP_BLOCK_NUM to either remove the redundant asserts or replace them
with independently-derived expectations that still validate the sparse-cache
geometry against config drift.
🪄 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
Run ID: 620632a4-c961-459d-a56f-e8c67ff73eb3
📒 Files selected for processing (22)
models/deepseek/v4/config.pymodels/deepseek/v4/decode_attention_csa.pymodels/deepseek/v4/decode_attention_hca.pymodels/deepseek/v4/decode_attention_swa.pymodels/deepseek/v4/decode_compressor_ratio128.pymodels/deepseek/v4/decode_compressor_ratio4.pymodels/deepseek/v4/decode_fwd.pymodels/deepseek/v4/decode_indexer.pymodels/deepseek/v4/decode_indexer_compressor.pymodels/deepseek/v4/decode_sparse_attn.pymodels/deepseek/v4/decode_sparse_attn_hca.pymodels/deepseek/v4/decode_sparse_attn_swa.pymodels/deepseek/v4/prefill_attention_csa.pymodels/deepseek/v4/prefill_attention_hca.pymodels/deepseek/v4/prefill_attention_swa.pymodels/deepseek/v4/prefill_compressor_ratio128.pymodels/deepseek/v4/prefill_compressor_ratio4.pymodels/deepseek/v4/prefill_fwd.pymodels/deepseek/v4/prefill_indexer.pymodels/deepseek/v4/prefill_indexer_compressor.pymodels/deepseek/v4/prefill_layer.pymodels/deepseek/v4/prefill_sparse_attn.py
high-cloud
force-pushed
the
fix/dsv4-cache-layout-constants
branch
from
e85d2ee to
50e316b
Compare
Define shared paged-cache block counts in config.py and use them across decode and prefill attention, compressor, sparse-attn, and indexer paths. Size prefill compressed KV and index KV pools from the same decode-batch physical layout while keeping per-table max-block widths explicit. Update L3 prefill/decode forward kernels to emit hidden states for the serving-side lm_head path, drop the stale lm_head_weight ABI, and raise the MoE recv-buffer floor for the current distributed layout.
high-cloud
force-pushed
the
fix/dsv4-cache-layout-constants
branch
from
50e316b to
b496207
Compare
2 participants
Summary
config.pyand reuse them across decode and prefill cache users.Related Issues
N/A