Add: DeepSeek V4 prefill support up to 8192 tokens - #893
Conversation
📝 WalkthroughWalkthroughThe PR adds configurable DeepSeek-V4 Flash prefill profiles, protocol-dependent MoE signaling, token-aware state mapping, and serving8k layer and finalization modules. ChangesDeepSeek V4 Flash prefill
Estimated code review effort: 5 (Critical) | ~120 minutes Possibly related issues
Possibly related PRs
Suggested labels: 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: 1
🧹 Nitpick comments (3)
models/deepseek/v4-flash/moe.py (1)
75-76: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd a bound assertion for the epoch slot index.
In the
epoch_isolatedprotocol,dispatchandcombineindex the signal windows atmoe_epoch - 1. The window hasMOE_SIGNAL_SLOTS = M.num_hidden_layersrows. Nothing here states that every caller's maximummoe_epochfits that row count. The packed graph drives epochs up toLAST_MOE_EPOCH, derived fromFWD_NUM_LAYERS. If those two layer counts ever diverge, the notify and wait offsets leave the window.Add an import-time assertion that ties the slot count to the maximum epoch the module supports.
🤖 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-flash/moe.py` around lines 75 - 76, Add an import-time assertion near MOE_SIGNAL_SLOTS and the epoch configuration to ensure the epoch-isolated slot count is at least LAST_MOE_EPOCH, preventing dispatch and combine indexing at moe_epoch - 1 from exceeding the signal window; preserve the existing non-isolated configuration.models/deepseek/v4-flash/prefill_layer_serving8k.py (1)
60-68: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winExtend the post-import check to
MOE_TOKENS.Line 53 mutates
config.MOE_TOKENSbefore importingmoe. That mutation only takes effect ifmoewas not already imported in this process. The check on Lines 60-68 re-validates the profile fingerprint and the signal protocol after import, but it does not inspect the token count. Amoemodule imported earlier with the decode token count therefore passes this guard while carrying the wrongT.Add
moe_module.T(or the equivalent token constant) to the stale-import check.🤖 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-flash/prefill_layer_serving8k.py` around lines 60 - 68, Extend the post-import validation condition in the prefill profile guard to also verify that moe_module.T matches the expected prefill MOE token count from config.MOE_TOKENS. Raise the existing RuntimeError when the imported module carries a stale decode token count, while preserving the current fingerprint and signal-protocol checks.models/deepseek/v4-flash/prefill_attention_csa.py (1)
610-611: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueBound the sparse-row estimate by the profile constant.
max_sparse_rowsusesIDX_TOPK, the raw model top-k. The kernel selects at mostPREFILL_MAX_COMPRESSEDcompressed rows, which equalsIDX_TOPKonly for the non-legacy profiles. For thelegacyprofile the effective bound is smaller, so this estimate is pessimistic. The comment on Line 1014 already describes the suffix in terms of the profile cap, so the two now use different constants.The current check still passes inside every profile's
RAW_TOKEN_CAP, so this is a consistency improvement rather than a defect.♻️ Proposed alignment with the profile bound
max_visible_cmp = (context_len + q_len) // COMPRESS_RATIO - max_sparse_rows = WIN + min(max_visible_cmp, IDX_TOPK) + max_sparse_rows = WIN + min(max_visible_cmp, PREFILL_MAX_COMPRESSED)🤖 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-flash/prefill_attention_csa.py` around lines 610 - 611, Update the max_sparse_rows calculation near max_visible_cmp to use the profile-specific PREFILL_MAX_COMPRESSED bound instead of raw IDX_TOPK, keeping the WIN offset and visible-row minimum logic unchanged. Align this estimate with the profile-cap suffix calculation described near the existing Line 1014 comment.
🤖 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-flash/prefill_layer_serving8k.py`:
- Around line 640-646: Update the L3 prefill layer flow so each call to
l3_prefill_layer uses a distinct per-layer moe_epoch, ensuring epoch_isolated
selects a separate signal row for every layer. Pass the prior layer’s epoch to
clear_moe_signal_epoch when clearing, and if epochs are reused, add the required
cross-rank barrier before clearing.
---
Nitpick comments:
In `@models/deepseek/v4-flash/moe.py`:
- Around line 75-76: Add an import-time assertion near MOE_SIGNAL_SLOTS and the
epoch configuration to ensure the epoch-isolated slot count is at least
LAST_MOE_EPOCH, preventing dispatch and combine indexing at moe_epoch - 1 from
exceeding the signal window; preserve the existing non-isolated configuration.
In `@models/deepseek/v4-flash/prefill_attention_csa.py`:
- Around line 610-611: Update the max_sparse_rows calculation near
max_visible_cmp to use the profile-specific PREFILL_MAX_COMPRESSED bound instead
of raw IDX_TOPK, keeping the WIN offset and visible-row minimum logic unchanged.
Align this estimate with the profile-cap suffix calculation described near the
existing Line 1014 comment.
In `@models/deepseek/v4-flash/prefill_layer_serving8k.py`:
- Around line 60-68: Extend the post-import validation condition in the prefill
profile guard to also verify that moe_module.T matches the expected prefill MOE
token count from config.MOE_TOKENS. Raise the existing RuntimeError when the
imported module carries a stale decode token count, while preserving the current
fingerprint and signal-protocol checks.
🪄 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: 4f823b03-47ee-4876-a80e-e95633555c9a
📒 Files selected for processing (9)
models/deepseek/v4-flash/config.pymodels/deepseek/v4-flash/moe.pymodels/deepseek/v4-flash/prefill_attention_csa.pymodels/deepseek/v4-flash/prefill_context_profile.pymodels/deepseek/v4-flash/prefill_finalize_serving8k.pymodels/deepseek/v4-flash/prefill_fwd.pymodels/deepseek/v4-flash/prefill_indexer.pymodels/deepseek/v4-flash/prefill_layer_serving8k.pymodels/deepseek/v4-flash/prefill_sparse_attn.py
b89e7d5 to
67f84c1
Compare
- prefill_fwd takes a whole request: the request-scoped tensors carry a
dynamic token extent, and the compiled body walks them in fixed-T tiles.
The 43-layer body still compiles for exactly T rows, and each tile binds its
activations to its own scope. The parameter list keeps the same names,
order, count and dtypes, and the tile count comes from the submitted extent.
- The paged caches stay on device across tiles. KV, compressed KV, indexer KV
and the three compressor states are InOut and addressed by slot mapping, so
each tile writes its own pages and later tiles read the accumulated history.
The body clears its MoE signal windows on exit. A tile past the end of a
rank's request runs one fully masked sentinel row.
- The MoE round declares the producer-to-waiter dependencies it relies on:
dispatch_push -> dispatch_wait, combine -> combine_wait and
dispatch_push -> combine_wait, and dispatch_gather is no longer
pre-resolved. A request now runs up to 64 tiles, so the number of
all-to-all rounds per dispatch grows with the prompt; without those edges a
blocking waiter can hold a core group while the local notifier still has
blocks parked in that core's pending slot, which deadlocks EP8 prefill.
- One tile reaches 8192 tokens of history. The Indexer scores up to 2048
compressed candidates, sorts them in full and emits the model-configured
top-k of 512; sparse attention covers 128 sliding-window rows plus those 512,
padded to 640.
- The packed CSA capacity bound compares against the Indexer's top-k rather
than the raw candidate count.
Callers size the request tensors to the prompt and submit one task:
x_hc = torch.empty(ranks, prompt_tokens, HC_MULT, D) # was prefill_seq rows
67f84c1 to
51fd55a
Compare
|
EP8 prefill on cards 0,2,4,6,8,10,12,14, both variants interleaved on the same No measurable cost.
Tile scaling at EP2, 128 / 256 / 512 tokens: 115627 / 223677 / 453552 us, i.e. Goldens on this branch rebased onto main: |
The dspark directory forked from deepseek_v4_flash_mtp before #893 and #953 landed, so it still carries the MoE handshake, split-K and indexer sort behaviour those two fixed. This ports their dspark-applicable parts; the stacked YaRN RoPE tables from #953 are not included, since they live in the decode_fwd / prefill_fwd routing dspark does not have yet. - Declare dispatch_push -> dispatch_wait and drop the pre-resolve on dispatch_gather, so a blocking waiter cannot hold a core group while the local notifier still has blocks parked in that core's pending slot. Without those edges EP8 prefill deadlocks once a request runs enough all-to-all rounds, and dspark runs EP16. - Publish the MoE combine arrivals from a combine_wait scope gated on the whole scatter grid and on dispatch_push, instead of folding one notify into each scatter block, so the wait expects moe_epoch rather than moe_epoch * N_LOCAL. The combine scatter, its wait and shared_routed lose allow_early_resolve, so the cross-rank handshake cannot reserve the AIV cores the scatter itself needs. moe.py is now identical to the mtp file apart from the EP naming and the 16-experts-per-rank split. - Replace the FP32 AtomicAdd accumulation in hc_pre's linear and RMS split-K paths with disjoint per-split partial buffers reduced in ascending K order, in both the fused and the separate implementation, and retire the zero-seed phase the atomics required. AtomicAdd sums the partials in task-completion order, which is neither fixed nor associative, so the same greedy request can diverge and then compound autoregressively over a long decode. The golden follows the same ascending order. - Add the block_len=1024 merge stage to the prefill indexer top-k when INDEXER_SCORE_CAP exceeds 256. dspark sizes the cap as 2 * T / COMPRESS_RATIO, so it is exactly 256 at the checked-in PREFILL_SEQ=512 and larger for any longer chunk; past 256 the 64/256 stages leave the score prefix partially ordered and the top-k silently picks the wrong keys. The CP score path already carries this stage for its 1024-wide cap. - Bound the packed sparse CSA rows by what the indexer actually emits, WIN + min(max_visible_cmp, IDX_TOPK), so build_tensor_specs stops refusing prompts that fit. The hc_pre reduction costs the determinism it buys: on a2a3, 100 rounds after 5 warmup, the 512-row prefill case goes from 167.0 to 199.5 us minimum over two runs, while the 128-row decode case is unchanged.
2 participants
dynamic token extent, and the compiled body walks them in fixed-T tiles.
The 43-layer body still compiles for exactly T rows, and each tile binds its
activations to its own scope. The parameter list keeps the same names,
order, count and dtypes, and the tile count comes from the submitted extent.
and the three compressor states are InOut and addressed by slot mapping, so
each tile writes its own pages and later tiles read the accumulated history.
The body clears its MoE signal windows on exit. A tile past the end of a
rank's request runs one fully masked sentinel row.
dispatch_push -> dispatch_wait, combine -> combine_wait and
dispatch_push -> combine_wait, and dispatch_gather is no longer
pre-resolved. A request now runs up to 64 tiles, so the number of
all-to-all rounds per dispatch grows with the prompt; without those edges a
blocking waiter can hold a core group while the local notifier still has
blocks parked in that core's pending slot, which deadlocks EP8 prefill.
compressed candidates, sorts them in full and emits the model-configured
top-k of 512; sparse attention covers 128 sliding-window rows plus those 512,
padded to 640.
than the raw candidate count.
Callers size the request tensors to the prompt and submit one task: