feat(deepseek-v4): Define dynamic prefill attention contract#793
feat(deepseek-v4): Define dynamic prefill attention contract#793hashiqiqixian wants to merge 2 commits into
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:
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 |
![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 a dynamic-token contract for DeepSeek V4 packed prefill, transitioning various attention, compressor, indexer, and sparse attention modules to use dynamic token shapes (T_DYN) instead of static dimensions. Feedback on the changes identifies an issue in hc_pre.py where pl.store is incorrectly passed a TensorType from pl.set_validshape instead of a Vec TileType. To fix this and meet hardware alignment constraints, it is recommended to allocate a same-core scratch tensor comb_pad_store to assemble the tiles and load them back as Vec tiles before storing.
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.
![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
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (3)
models/deepseek/v4/prefill_attention_csa.py (1)
263-292: 🩺 Stability & Availability | 🔴 Critical | ⚡ Quick winGuard the tail-tile assemblies.
For non-aligned token counts,
t_idxexceeds the exact[num_tokens, ...]outputs on the final tile, but Lines 291-292 still assemble both rows out of bounds.Proposed fix
if t_idx < num_tokens: abs_pos = pl.read(position_ids, [t_idx]) ... - swa_indices = pl.assemble(swa_indices, swa_row, [t_idx, 0]) - cmp_indices = pl.assemble(cmp_indices, cmp_row, [t_idx, 0]) + swa_indices = pl.assemble(swa_indices, swa_row, [t_idx, 0]) + cmp_indices = pl.assemble(cmp_indices, cmp_row, [t_idx, 0])🤖 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_csa.py` around lines 263 - 292, Guard the swa_indices and cmp_indices assemblies in the top-k tile loop so they run only when t_idx is less than num_tokens. Keep the existing row initialization and valid-token population unchanged, and avoid assembling either row for padded tail-tile positions.models/deepseek/v4/prefill_indexer.py (1)
251-259: 🗄️ Data Integrity & Integration | 🔴 Critical | ⚡ Quick winPass the dynamic input to the child compressor.
Line 252 passes padded
[T, D]tensorx, soprefill_indexer_compressorderivesnum_tokens == Tand can overrun the shorter metadata tensors. Passx_in; the child already performs its own static padding.Proposed fix
prefill_indexer_compressor( - x, + x_in,🤖 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_indexer.py` around lines 251 - 259, Update the prefill_indexer_compressor call to pass the dynamic input tensor x_in instead of the padded tensor x, preserving the child compressor’s existing static-padding behavior and ensuring its token count matches the metadata tensors.models/deepseek/v4/prefill_indexer_compressor.py (1)
463-469: 🎯 Functional Correctness | 🔴 Critical | ⚡ Quick winDedent the scan-loop body after removing the guard.
Lines 465-469 remain indented under the deleted conditional, making the module fail Python parsing.
Proposed fix
num_tokens = pl.tensor.dim(x, 0) for scan_w in pl.range(num_tokens): - scan_slot_raw = pl.read(idx_slot_mapping, [scan_w]) - if scan_slot_raw >= 0: - if write_seen == kv_i: - src_row_raw = scan_slot_raw - write_seen = write_seen + 1 + scan_slot_raw = pl.read(idx_slot_mapping, [scan_w]) + if scan_slot_raw >= 0: + if write_seen == kv_i: + src_row_raw = scan_slot_raw + write_seen = write_seen + 1🤖 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_indexer_compressor.py` around lines 463 - 469, Dedent the scan-loop body in the token scan logic around num_tokens so scan_slot_raw handling and write_seen updates are directly nested under the pl.range loop after the removed guard. Ensure the resulting indentation is valid Python and preserves the existing write_seen and src_row_raw behavior.
🧹 Nitpick comments (1)
tests/golden/test_deepseek_v4_prefill_dynamic_source.py (1)
19-21: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueCache AST parsing and reuse
_tree.This helper is called repeatedly across multiple tests (e.g., inside
_function), which redundantly reads and parses the same files multiple times. Additionally, line 61 manually parses the AST instead of using_tree.Consider caching the ASTs to improve test performance and replacing the inline parsing at line 61 with
tree = _tree(filename).♻️ Proposed refactor
+import functools + +@functools.lru_cache(maxsize=None) def _tree(filename: str) -> ast.Module: return ast.parse((MODEL / filename).read_text(encoding="utf-8"))And update line 61:
- tree = ast.parse((MODEL / filename).read_text(encoding="utf-8")) + tree = _tree(filename)🤖 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/golden/test_deepseek_v4_prefill_dynamic_source.py` around lines 19 - 21, Cache the parsed AST result in the _tree helper so repeated requests for the same filename reuse the existing module instead of rereading and reparsing it. Update the inline AST parsing near the affected test flow to assign tree via _tree(filename), preserving the current parsing behavior while centralizing it in _tree.
🤖 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/prefill_indexer.py`:
- Around line 566-577: Update the standalone CLI entrypoints so each builder
receives the requested token count: add a --num-tokens argument and forward
args.num_tokens when calling build_tensor_specs in
models/deepseek/v4/prefill_indexer.py (566-577) and
models/deepseek/v4/prefill_indexer_compressor.py (610-622); pass args.num_tokens
alongside args.compress_ratio in models/deepseek/v4/prefill_sparse_attn.py
(638-654).
- Around line 530-536: Update the topk copy loop around cmp_topk_indices to
initialize each topk_idxs row to -1, then copy only IDX_TOPK columns from the
IDX_TOPK-wide source tensor. Do not slice cmp_topk_indices through
INDEXER_SCORE_CAP; preserve the existing token bounds and tiling behavior.
- Around line 404-408: Update the remaining golden calculations after num_tokens
is derived to use num_tokens instead of the stale T extent, including view
shapes and the reshape involving IDX_N_HEADS. Ensure all intermediate tensors
are sized to the actual token count so partial-token fixtures compute correctly.
---
Outside diff comments:
In `@models/deepseek/v4/prefill_attention_csa.py`:
- Around line 263-292: Guard the swa_indices and cmp_indices assemblies in the
top-k tile loop so they run only when t_idx is less than num_tokens. Keep the
existing row initialization and valid-token population unchanged, and avoid
assembling either row for padded tail-tile positions.
In `@models/deepseek/v4/prefill_indexer_compressor.py`:
- Around line 463-469: Dedent the scan-loop body in the token scan logic around
num_tokens so scan_slot_raw handling and write_seen updates are directly nested
under the pl.range loop after the removed guard. Ensure the resulting
indentation is valid Python and preserves the existing write_seen and
src_row_raw behavior.
In `@models/deepseek/v4/prefill_indexer.py`:
- Around line 251-259: Update the prefill_indexer_compressor call to pass the
dynamic input tensor x_in instead of the padded tensor x, preserving the child
compressor’s existing static-padding behavior and ensuring its token count
matches the metadata tensors.
---
Nitpick comments:
In `@tests/golden/test_deepseek_v4_prefill_dynamic_source.py`:
- Around line 19-21: Cache the parsed AST result in the _tree helper so repeated
requests for the same filename reuse the existing module instead of rereading
and reparsing it. Update the inline AST parsing near the affected test flow to
assign tree via _tree(filename), preserving the current parsing behavior while
centralizing it in _tree.
🪄 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: 792a0f7d-4390-47bd-a37f-ef0865b00ad9
📒 Files selected for processing (14)
models/deepseek/v4/dynamic_shapes.pymodels/deepseek/v4/hc_post.pymodels/deepseek/v4/hc_pre.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_indexer.pymodels/deepseek/v4/prefill_indexer_compressor.pymodels/deepseek/v4/prefill_sparse_attn.pymodels/deepseek/v4/qkv_proj_rope.pymodels/deepseek/v4/rmsnorm.pytests/golden/test_deepseek_v4_prefill_dynamic_source.py
0a011c8 to
41f0ee5
Compare
5a57853 to
dadfe8f
Compare
329f905 to
1944a03
Compare
1944a03 to
82f7542
Compare
1 participant
Summary
TOKENS_DYNcontract for the DeepSeek V4 prefill Attention path and bind token-aligned inputs, outputs, positions, and slot mappings at JIT entry points.num_tokensargument from the Attention leaf APIs and propagate dynamically shaped tensors through the SWA, HCA, and CSA dependency chains.validshapefor non-tile-aligned tail rows.num_tokensonly at the three top-level Attention composition entry points as a temporary compatibility boundary for existing fixed-Tcallers. Removing that boundary, converting MoE, and integrating the full packedprefill_layer/prefill_fwdpath are intentionally deferred to the follow-up PR.python -m pytest tests/golden/test_deepseek_v4_prefill_dynamic_source.py -q— 8 passedgit diff --check origin/main...HEADRelated Issues
Part of #766
Related to #410