Add: dynamic decode batch across the DeepSeek-V4 Flash dspark tree - #910
Merged
zhangqi-chen merged 5 commits intoAug 7, 2026
Merged
zhangqi-chen merged 5 commits into
zhangqi-chen merged 5 commits into
Conversation
Serve any vLLM ACL-graph decode bucket from one compiled program. The A3
single-node deployment (TP4 DP4, --max-num-seqs 64, MTP=1, FULL_DECODE_ONLY,
FLASHCOMM1) captures token counts [4, 8, 16, 24, ..., 128], so the runtime
contract is B a multiple of 4 up to 64, i.e. T = B*S a multiple of 8 up to 128.
DECODE_BATCH stays the compile-time upper bound that sizes all scratch.
Every kernel from rope_interleave up to decode_swa / decode_hca / decode_csa
now derives its token count with pl.tensor.dim and drives loop bounds from it;
each test entry gains -b/--batch and its fixture is parameterized.
Three classes of latent bug surfaced, all invisible at B = 64:
- Floor-divided cube-tile grids dropped the tail token block. qkv_proj_rope and
decode_compressor_ratio128 computed the block count as t // TILE, so at
T = 40 the last 8 tokens were never projected. Now ceil.
- A matmul whose M tile is the static token count is only valid while
bs == T_PAD. decode_indexer, and the proj_a of all three sparse_attn
kernels, asked for a tile taller than the source tensor (or than the written
region of a T_PAD-strided scratch) and silently produced zeros. Those
matmuls are row-blocked at the 16-row cube floor.
- A rank-increasing pl.reshape feeding an @pl.jit.inline parameter fails
metadata inference once the axis is dynamic. The compressors and the indexer
therefore take token-major flat tensors, matching how vLLM hands over slot
mappings; b_dim now comes from a per-request block table.
Also fixes an orchestration race the batch work exposed: decode_indexer's
`score` is written by one scope and read back by topk, but pypto emits
add_output for a pl.Out parameter, so the RAW edge was missing. It is pl.InOut
now.
Validated on a2a3 at B = 64 and at least one value of the B = 4 (mod 8) class
(12/20/32/60), which is the class that exercises the partial 16-row tile.
DEPLOYMENT NOTE: the T = 4 capture bucket must be pinned out, e.g.
--compilation-config '{"cudagraph_mode":"FULL_DECODE_ONLY",
"cudagraph_capture_sizes":[8,16,...,128]}'. Without it a 1-2 request batch
lands on a bucket the kernels do not support and returns uninitialized memory
rather than failing.
Comments state what, not why: drop the "divides by a compile-time constant" rationale from the seven spmd index decodes, collapse the two-line pad-zero and row-block notes to one line each, and reword the flat-ABI notes to name the constraint instead of retelling the diagnosis. Header groups follow the model-config / tiling convention: decode_indexer had the dynamic-variable group under the `# model config` label, decode_sparse_attn_swa ended up with two `# Dynamic shape variables.` labels, and rope_interleave had no group labels at all. Split the one over-long statement introduced by the pass: the hca rope_cs gather stores are now one op per line. Also drop seven pre-existing dead locals the pass surfaced -- five stale reads in the goldens plus a duplicated pl.tensor.dim in the swa entry. The two `as qk_tid` captures stay: pl.spmd only accepts `deps=` in the capture form, so dropping the name fails to compile. They are renamed `_qk_tid` instead.
DECODE_BATCH is the per-DP-rank request count. TP is not implemented, so a
single card carries the whole rank and every decode kernel sized its batch
axis 4x too large. Use DECODE_BATCH // TP for the compile-time bound in the
fourteen non-MoE decode kernels; expert_routed keeps DECODE_BATCH.
decode_mtp_verify took T from DECODE_TOKENS rather than B * S, which no longer
agreed with the per-card B; it derives T from B now and the assert it guarded
is redundant.
B_MAX is 16, so the dynamic batch contract narrows to B in {4, 8, 12, 16}
(T in {8, 16, 24, 32}).
rmsnorm, hc_pre, hc_post, qkv_proj_rope and mtp_projection serve both paths from one T_DYN kernel, so their decode test point and their decode-side tile asserts still read the per-DP-rank DECODE_BATCH. Match the fused kernels and use DECODE_BATCH // TP; the prefill mode is untouched.
|
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:
📝 WalkthroughWalkthroughThe DeepSeek V4 decode paths now use tensor-parallel-adjusted batch capacity, flattened token-major tensors, runtime dynamic dimensions, padded token workspaces, and configurable batch fixtures. Decode CLIs validate runtime batch values. ChangesDeepSeek V4 dynamic decode
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: 6
🧹 Nitpick comments (1)
models/deepseek_v4_flash_dspark/decode_indexer.py (1)
748-761: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winThe
--batchdefault is not checked against its own validation rule. All seven harnesses default-b/--batchtoB = DECODE_BATCH // TPand then reject any value below 4 or not a multiple of 4. If a configuration yieldsB < 4orB % 4 != 0, running the harness with no-bflag fails with a parser error. The validation logic and the help text are also copied verbatim across all seven files.
models/deepseek_v4_flash_dspark/decode_indexer.py#L748-L761: extract the argument definition and the range check into one shared helper inmodels/deepseek_v4_flash_dspark/utils.py, and assert at import time thatBsatisfies the rule.models/deepseek_v4_flash_dspark/decode_indexer_compressor.py#L642-L658: replace the local definition and check with the shared helper.models/deepseek_v4_flash_dspark/decode_csa.py#L920-L939: replace the local definition and check with the shared helper.models/deepseek_v4_flash_dspark/decode_hca.py#L691-L708: replace the local definition and check with the shared helper.models/deepseek_v4_flash_dspark/decode_sparse_attn_csa.py#L807-L827: replace the local definition and check with the shared helper.models/deepseek_v4_flash_dspark/decode_sparse_attn_hca.py#L806-L827: replace the local definition and check with the shared helper.models/deepseek_v4_flash_dspark/decode_sparse_attn_swa.py#L684-L701: replace the local definition and check with the shared helper.🤖 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_dspark/decode_indexer.py` around lines 748 - 761, The --batch default must satisfy the same validation rule applied to user input. In models/deepseek_v4_flash_dspark/utils.py, add a shared helper that defines the -b/--batch argument and performs the range and multiple-of-4 check, and assert at import time that B is at least 4 and divisible by 4. Replace the local argument definition and validation in decode_indexer.py (748-761), decode_indexer_compressor.py (642-658), decode_csa.py (920-939), decode_hca.py (691-708), decode_sparse_attn_csa.py (807-827), decode_sparse_attn_hca.py (806-827), and decode_sparse_attn_swa.py (684-701) with this helper.
🤖 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_dspark/config.py`:
- Line 242: Clarify the DECODE_BATCH contract in the configuration and its
consumers: either define separate total and per-card values, or update local
decode-kernel calculations to derive the intended per-card batch consistently,
including the MoE path using RECV_MAX. Ensure comments and all usages agree on
whether DECODE_BATCH is DP-sized or local per-card.
In `@models/deepseek_v4_flash_dspark/decode_compressor_ratio128.py`:
- Around line 109-110: Update the cos and sin tensor shape annotations in
compressor_ratio128() to use the static B capacity instead of B_DYN, matching
decode_compressor_ratio4.py and the buffers written by rope_interleave().
In `@models/deepseek_v4_flash_dspark/decode_hca.py`:
- Around line 149-152: Update the block-count calculations near t_dim, b_dim,
and the topk_blocks/wb_blocks assignments to use ceiling division for both
HCA_TOPK_TOKEN_TILE and HCA_WB_TOKEN_TILE. Preserve the existing topk_t < t_dim
bounds guard and add the equivalent write_t < t_dim guard in the writeback loop
so tail tokens are initialized and written safely.
In `@models/deepseek_v4_flash_dspark/decode_indexer.py`:
- Around line 389-390: Update the cos and sin parameter annotations in the
relevant decode-indexer function to use B_DYN instead of B, matching their
bind_dynamic declarations and the batch-shaped fixture output. Keep the existing
second dimension ROPE_HEAD_DIM // 2 and dynamic bindings unchanged.
In `@models/deepseek_v4_flash_dspark/decode_sparse_attn_hca.py`:
- Around line 142-156: Token-tile grids currently floor-divide t_dim and drop
tail tokens; update
models/deepseek_v4_flash_dspark/decode_sparse_attn_hca.py:142-156 to use ceil
grids for valid_blocks and act_t_blks and guard per-token writes. Apply the same
fix in models/deepseek_v4_flash_dspark/decode_hca.py:149-152 for topk_blocks and
wb_blocks, adding the write_t < t_dim writeback guard; in
models/deepseek_v4_flash_dspark/decode_csa.py:169-171 for wb_blocks and its
writeback loop at line 240; and in
models/deepseek_v4_flash_dspark/decode_sparse_attn_swa.py:394-398 for act_t_blks
and proj_b_act stores. Preserve existing topk guards and ensure sparse_bias rows
and tail-token activations are initialized and written.
In `@models/deepseek_v4_flash_dspark/qkv_proj_rope.py`:
- Line 119: Update workspace allocations around t_matmul in
models/deepseek_v4_flash_dspark/qkv_proj_rope.py: use the static decode capacity
for create_tensor shapes instead of runtime-derived t_dim or t_matmul, while
restricting computation through valid_shape or active-token loop bounds. Apply
the same static-capacity workspace change to the affected create_tensor calls in
models/deepseek_v4_flash_dspark/decode_swa.py lines 113-140.
---
Nitpick comments:
In `@models/deepseek_v4_flash_dspark/decode_indexer.py`:
- Around line 748-761: The --batch default must satisfy the same validation rule
applied to user input. In models/deepseek_v4_flash_dspark/utils.py, add a shared
helper that defines the -b/--batch argument and performs the range and
multiple-of-4 check, and assert at import time that B is at least 4 and
divisible by 4. Replace the local argument definition and validation in
decode_indexer.py (748-761), decode_indexer_compressor.py (642-658),
decode_csa.py (920-939), decode_hca.py (691-708), decode_sparse_attn_csa.py
(807-827), decode_sparse_attn_hca.py (806-827), and decode_sparse_attn_swa.py
(684-701) with this helper.
🪄 Autofix
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: 573a1996-ed67-4d77-ae06-ccc865808012
📒 Files selected for processing (20)
models/deepseek_v4_flash_dspark/config.pymodels/deepseek_v4_flash_dspark/decode_compressor_ratio128.pymodels/deepseek_v4_flash_dspark/decode_compressor_ratio4.pymodels/deepseek_v4_flash_dspark/decode_csa.pymodels/deepseek_v4_flash_dspark/decode_hca.pymodels/deepseek_v4_flash_dspark/decode_indexer.pymodels/deepseek_v4_flash_dspark/decode_indexer_compressor.pymodels/deepseek_v4_flash_dspark/decode_metadata.pymodels/deepseek_v4_flash_dspark/decode_mtp_verify.pymodels/deepseek_v4_flash_dspark/decode_sparse_attn_csa.pymodels/deepseek_v4_flash_dspark/decode_sparse_attn_hca.pymodels/deepseek_v4_flash_dspark/decode_sparse_attn_swa.pymodels/deepseek_v4_flash_dspark/decode_swa.pymodels/deepseek_v4_flash_dspark/hc_head.pymodels/deepseek_v4_flash_dspark/hc_post.pymodels/deepseek_v4_flash_dspark/hc_pre.pymodels/deepseek_v4_flash_dspark/mtp_projection.pymodels/deepseek_v4_flash_dspark/qkv_proj_rope.pymodels/deepseek_v4_flash_dspark/rmsnorm.pymodels/deepseek_v4_flash_dspark/rope_interleave.py
compressor_ratio128 declared its interleaved cos/sin as [B_DYN, ...], but those buffers are the B_MAX-sized scratch rope_interleave writes, which is how compressor_ratio4 and indexer_compressor already declare them. indexer_test declared its half-width cos/sin with the static B while binding dimension 0 to B_DYN, so the annotation disagreed with both the bind and the fixture.
This was referenced Aug 7, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Serve any decode batch from one compiled program across the dspark
decode tree, from rope_interleave up to decode_swa / decode_hca /
decode_csa. Each kernel derives its token count with pl.tensor.dim and
drives every loop bound from it; each test entry gains -b/--batch and
its fixture is parameterized.
Size the decode kernels for one card rather than one DP rank.
DECODE_BATCH is the per-DP-rank request count and TP is not
implemented, so a single card carries the whole rank; the fourteen
non-MoE decode kernels and the shared decode/prefill leaves now bound
their batch axis with DECODE_BATCH // TP. expert_routed keeps
DECODE_BATCH. The compile-time bound is 16, so the runtime contract is
B in {4, 8, 12, 16}, i.e. T = B * S in {8, 16, 24, 32}.
Ceil the cube-tile grids that previously floored. qkv_proj_rope and
decode_compressor_ratio128 computed their block count as t // TILE, so
a token count that was not a whole multiple of the 16-row tile left the
last block unprojected.
Row-block every matmul whose M tile was the static token count.
decode_indexer and the proj_a of all three sparse_attn kernels asked
for a tile taller than the source tensor, or taller than the written
region of a T_PAD-strided scratch, and produced zeros. Those matmuls
fan ceil(t / 16) blocks at the 16-row cube floor.
Take token-major flat tensors in the compressors and the indexer. A
rank-increasing pl.reshape feeding an @pl.jit.inline parameter fails
metadata inference once the axis is dynamic, so the token axis stays
flat end to end and b_dim comes from a per-request block table. This
also matches the order vLLM hands over slot mappings.
Make decode_indexer's score parameter pl.InOut. It is written by one
scope and read back by the topk scope, and pypto emits add_output for
a pl.Out parameter, so the orchestration missed the read-after-write
edge.
Deployment note: the kernels do not support a decode batch below 4, so
the 4-token ACL-graph capture bucket must be pinned out with
cudagraph_capture_sizes. On that bucket the tile grids run zero blocks
and return uninitialized memory rather than failing.