fix(cache): support attention-DP for Kimi-K3 by deriving the MLA packing from the KDA state size - #1152
Open
dongjiyingdjy wants to merge 3 commits into
Open
fix(cache): support attention-DP for Kimi-K3 by deriving the MLA packing from the KDA state size#1152dongjiyingdjy wants to merge 3 commits into
dongjiyingdjy wants to merge 3 commits into
Conversation
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: a6ee91b3e6
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
Comment on lines
+270
to
+271
| mla_packing = max(1, -(-linear_plane_bytes // mla_page_bytes)) | ||
| linear_packing = max(1, mla_packing * mla_page_bytes // linear_plane_bytes) |
There was a problem hiding this comment.
With a BF16 MLA cache and attn_tp_size=32—both accepted by _kda_shapes—the 203,520-byte KDA state and 147,456-byte latent page make this choose MLA/state packing 2/1. Across 24 MLA planes but only 23 state slots, pack() then computes roughly 51% padding for each linear-attention group and rejects startup against max_padding_fraction=0.25. The previous 12/8 packing passed, and increasing the new MLA packing to 3 permits 3/2 with about 13% padding, so selecting only the smallest packing introduces a boot regression for this supported configuration.
Useful? React with 👍 / 👎.
Contributor
|
Hi @dongjiyingdjy please fix the conflicts thanks |
dongjiyingdjy
added a commit
that referenced
this pull request
Aug 21, 2026
Two blockers: restore --dist-init-addr in the DP config (dp>1 on bare
metal hard-raises without it; the earlier review removal only applies
to the Slurm-launched TP configs) and guard the conditional summary
keys in collect_outputs.py ("Decoded Tok/Iter" is only emitted with
speculative decoding active, which this matrix runs without; a missing
key crashed the whole CSV).
Also: declare the second DP prerequisite (#1185 merge_state head
tiling — attn tp=1 runs 96 heads/rank and the old kernel rejects the
launch on any cross-chunk merge) next to #1152 with the merge-order
note; raise both startup timers to 3600s (the harness timeout alone is
ineffective because ts serve's own engine_startup_timeout defaults to
1800s); spell out the builder's real --model-path flag instead of
relying on argparse prefix abbreviation; extend the fatal-error grep
with K3's known loud failures (ValueError cache-plan rejections,
NoKernelFoundError); fail fast if the dataset builds fewer than the 70
conversations the sweep consumes; correct the warmup wording in the
README.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: dongjiyingdjy <87510204+dongjiyingdjy@users.noreply.github.com>
dongjiyingdjy
added a commit
that referenced
this pull request
Aug 21, 2026
Fold DSpark into all three parallelism configs (matching the kimi_k2.5 convention of spec-on perf rows) and retire the separate _dspark config. All rows: util 0.92, kvstore on (DSPARK+KVStore validated on-machine — byte-identical temperature-0 answers vs the kvstore-off control, acceptance inside the run-to-run noise band, and 115 forced retracts restoring 30K+ cached-token prefixes through the L2 tier with zero exceptions), prefill graph enabled per review direction (the README carries the measured ~99.5% memory note for the TP8 rows). The DP row keeps its explicit chunk knobs and now carries DSpark too; the DP x DSpark combination is marked unvalidated in the config comments. The README is rewritten in the kimi_k2.5 operator-manual style: usage, workload sizing (68.2K worst-case final prompt -> max-model-len 80000), config table, the #1152 + #1185 merge order, the evalscope pin caveat, and the parallelism-verification tip. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: dongjiyingdjy <87510204+dongjiyingdjy@users.noreply.github.com>
The recipe pinned the MLA packing to a constant 12: the smallest plane
width (12 x 73,728B latent pages) that covers one per-layer KDA state at
attn tp=8. Under attention-DP (attn tp < 8) the per-rank state grows
8/tp-fold, outgrows the pinned plane, and the planner widens the plane to
fit it -- breaking the latent field's exact page stride and failing every
boot on the exact-page-stride check.
Derive the packing instead as the smallest count whose plane covers one
per-layer state page, ceil(state_bytes / latent_page_bytes): tp=8 keeps
its historical 12 and byte-identical layout, tp=16 halves the plane and
the parent with it, and attention-DP layouts (tp 4/2/1 -> packing
23/45/89) become plannable. The state-page padding stays within 13.4%
across every attn tp dividing the head count, in both cache dtypes.
Verified end to end on 8xB300 with nvidia/Kimi-K3-NVFP4 at
data-parallel-size 8 / ep-size 8: parent_bytes=157,483,008, packing
{full_attention: 89, linear_attention_*: 1}, 1.4M-token capacity.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: dongjiyingdjy <87510204+dongjiyingdjy@users.noreply.github.com>
dongjiyingdjy
force-pushed
the
fix/kimi-k3-dp-mla-packing
branch
from
August 25, 2026 13:35
a6ee91b to
5b823a0
Compare
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
2 participants
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.
Motivation: attention-DP for Kimi-K3
Kimi-K3 has never been able to boot under attention data parallelism: every
attn_tp < 8topology (e.g.--data-parallel-size 8 --ep-size 8on one node) fails cache planning withThis PR removes that blocker at the cache-recipe layer and is verified end to end serving nvidia/Kimi-K3-NVFP4 at DP8 x EP8 on 8xB300 (boot, correct concurrent generations across DP ranks, 1.4M-token capacity).
Root cause
The recipe pins the MLA packing to a constant 12. That constant is actually the value of a formula at attn tp=8: the smallest plane width (in 73,728 B latent pages) that covers one per-layer KDA state page:
Under attention-DP the state is not sharded (all 96 heads per rank), so it outgrows the pinned 884,736 B plane. The planner widens the plane to fit the state, which breaks the latent field's
exact_page_stridecontract (the MLA kernel indexes pages by an implicit payload-sized stride) and boot fails.Fix
Compute the packing as that formula instead of its tp=8 value:
tp=8 keeps today's layout byte-for-byte. tp=16 legitimately shrinks: half the state needs half the plane, so the parent halves at identical density (the two tests pinning the old tp16 values are updated accordingly). The latent page stride stays exactly 73,728 in every configuration -- now asserted by a dedicated regression test, since tp=8 is the one point where floor and ceil coincide and a formula regression would otherwise be invisible to the suite.
Verification
test/runtime/test_kimi_k3_cache_spec.py,test_kimi_k3_config.py: 37 passed. New tests: attention-DP layouts (tp 4/2/1), latent-stride exactness across tp 1..16, a bf16 cache-dtype case (the formula is a ratio of dtype-dependent byte counts; the suite was fp8-only).pack()sweep over every attn tp dividing 96, both cache dtypes.nvidia/Kimi-K3-NVFP4,--data-parallel-size 8 --ep-size 8 --moe-tp-size 1:Cache profile: parent_bytes=157483008, ... groups={'full_attention': 89, 'linear_attention_0': 1, ...}, warmup + concurrent requests correct on all DP ranks, token_capacity 1,446,656.Notes for DP deployments
--dense-tp-size(leave it at its defaultattn_tp): K3's dense MLP is not CommManager-routed and has a zero-token early return, so a dense group spanning DP ranks deadlocks (token-holding ranks enter the dense all-reduce while idle ranks skip ahead into the MoEtoken_all_gather). Observed and py-spy-verified; a startup guard would make this a 1-second error and could come as a follow-up.merge_statekernel fails withinvalid argumentat attn tp=1 (64 MLA heads/rank), hit on chunked prefill / prefix-cache merges.🤖 Generated with Claude Code