fix(vllm): make Kimi-K3 AttnRes hidden-state capture invariant to the PP partition - #167
Merged
Merged
Conversation
A capture id at a non-last stage's end_layer needs that layer's mixture weights, which live on the next rank. #50487 falls back to the running prefix there, a different feature; rejecting the id instead, as this patch previously did, does not survive mostly-PP sharding. Keep a copy of the two hidden-size vectors on the stage that taps them. Separately, the handoff folds the delta into the prefix, which drops the first layer of every non-first stage onto the Triton fallback and shifts the arithmetic for everything after it. Feed that layer a zero delta. Both verified bit-identical against PP=1 on GB300 at every split point. Signed-off-by: torchspec-bot <262938024+torchspec-bot@users.noreply.github.com>
test_pp_handoff.py fails without the zero-delta fix. test_eagle3.py needed two fixture updates from the #50487 backport, and now also asserts the tap passes no delta and no block write, which is what makes it a pure read. Signed-off-by: torchspec-bot <262938024+torchspec-bot@users.noreply.github.com>
Why the tap is a second kernel call, what the two PP bugs were, and what capture and the handoff cost. Signed-off-by: torchspec-bot <262938024+torchspec-bot@users.noreply.github.com>
torchspec-bot
force-pushed
the
fix/k3-attnres-pp-invariant
branch
from
August 10, 2026 02:58
d583a7f to
a20ec96
Compare
yubofredwang
approved these changes
Aug 10, 2026
yubofredwang
left a comment
yubofredwang
left a comment
Collaborator
There was a problem hiding this comment.
LGTM! thx for the fix.
| @@ -0,0 +1,251 @@ | |||
| # Capturing Kimi-K3 Hidden States under AttnRes | |||
Collaborator
There was a problem hiding this comment.
I think we can make this into a single recipe doc in the future.
| @@ -28,6 +28,31 @@ logs the active mode once at configuration time: | |||
| Once the base image carries PR #50487, drop this patch but keep exporting | |||
Collaborator
There was a problem hiding this comment.
confirm this patch builds correctly
torchspec-bot
force-pushed
the
fix/k3-attnres-pp-invariant
branch
from
August 10, 2026 03:03
a20ec96 to
d20918d
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.
Two bugs made the exported Kimi-K3 aux hidden states depend on where the pipeline was cut. Both are in the patched vLLM image, both are fixed here, and both are verified bit-identical against a PP=1 run on GB300.
Boundary tap
Under AttnRes the feature for capture id
kis a softmax mixture computed with layerk'sself_attention_res_norm/_proj. Whenkis a non-last stage'send_layer, those weights live on the next rank. Upstream vllm falls back to the running prefix there, which is a different feature entirely — cosine 0.05 to 0.60 against the real mixture. Rejecting the id instead, which is what this patch did before, does not survive mostly-PP sharding: for the 93-layer stage-1 recipe every intermediate id is a stage end from pp=40 on.Everything else the tap needs is already local, so a non-last AttnRes stage now keeps a copy of those two hidden-size vectors (~28KB, unsharded) and computes the mixture itself. A checkpoint that does not supply them fails at configuration time rather than exporting NaNs.
Kernel dispatch after a handoff
Found while measuring the above.
attn_resonly takes its fused native path when adeltais supplied, and the PP handoff folds the delta into the prefix. The first layer of every non-first stage was therefore running the Triton fallback, which differs from the native op on ~0.008% of elements and compounds through the rest of the stack — so the arithmetic, not just the boundary tap, depended on the partition. An explicit zero delta restores the dispatch, and is skipped for block-write layers, which never reach that path.Verification
GB300, real kernel, 32 layers / block size 8: the boundary tap is bit-identical to the PP=1 capture at every split point, and end-to-end bit-identity across the whole two-stage run goes from 4 of 31 split points to 31 of 31.
tests/vllm_k3_tests.patchis new.test_pp_handoff.pyfails without the zero-delta fix;test_eagle3.pypicks up two fixture updates from the #50487 backport and now also asserts the tap passes no delta and no block write, which is what keeps it a pure read of state the next stage consumes.Background and derivations are in
docs/kimi_k3_attn_res_capture.md.Test plan
tests/models/kimi_k3/test_pp_handoff.py+test_eagle3.py: 7 passed in the pinned imageruff check/ruff format --checkclean on the patched filesattn_reskernel