fix(npu): split packed MTP HiCache transfers by pool - #35013
Draft
momaekar1 wants to merge 22 commits into
Draft
Conversation
momaekar1
force-pushed
the
fix/npu-hicache-mtp-packed-transfer
branch
from
August 17, 2026 02:49
117edb4 to
98ae7de
Compare
7 tasks
momaekar1
force-pushed
the
fix/npu-hicache-mtp-packed-transfer
branch
from
August 17, 2026 09:04
98ae7de to
06ad15e
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
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.
Important
For standalone Ascend HiCache validation, this PR branch bundles the complete
commit history of #32500 through
#32500@9acc6b733. Users can fork this branchand run it directly without fetching another PR. For upstream review and
merge, #32500 should still land first; GitHub will then remove the shared
commits from this PR's diff. #34560 is already merged into
main. End-to-endQwen3.6 validation also uses the checkpoint semantic fix in #35086, but that
PR is not included in this branch. Recommended merge order: #32500, #35086,
then this PR.
Motivation
After #34560, the HiCache host K/V layer dimension includes both the target
model layers and the NEXTN/MTP draft layers. #32500 adds the Ascend fused
transfer_kv_dim_exchangepath forpage_first_direct, but that path invokesthe operator once with the target device pool and the full packed host buffer.
On Ascend, the target and draft K/V layers are separate device pools. A single
invocation therefore pairs a target-only device buffer with a
target + drafthost-layer range, so the draft K/V layers are not transferredthrough the matching device pool.
CUDA does not have this mismatch. Its existing path packs target and draft
per-layer pointers into one pointer array, and the per-layer transfer kernels
map each host layer to the corresponding device pointer. The Ascend fused
operator accepts one contiguous device pool per invocation, so each pool must
be transferred separately.
Minimal reproduction
Apply #32500 on top of a revision that contains #34560, then launch Qwen3.6-27B
on Ascend with the following relevant options:
from L1, then replay the original prompt.
/flush_cache, then replaythe original prompt.
return_cached_tokens_detailsand inspect the target/draft transferat the first
kernel_ascendH2D or D2H operation.Without this change, the fused operation receives the full packed host layer
range with only the target device pool. The draft pool is not transferred by a
matching invocation.
Changes
device pool.
transfer_kv_dim_exchangeonce per pool for both H2D and D2H.slicing.
9acc6b733so this fork branch is directlyrunnable for Ascend HiCache validation.
Focused validation:
upstream/main:2 passed#32500@9acc6b733:3 passed, 17 deselected8 passedL2/L3 replay validation
Validation stack: latest
mainwith #34560,#32500@9acc6b733, #35086, andthis fix. Model: Qwen3.6-27B BF16, TP2, Ascend NPU, NEXTN enabled, 4 prompts
with 3000 input tokens and 512 output tokens.
host=2944, device=0, storage=0for 4/4 promptsstorage=2944, device=0, host=0for 4/4 promptsThe L2 and L3 replay output IDs were identical for all four prompts:
Their speculative statistics were also identical:
62.8169%2.88451338/2130Cold generation and cache replay were not token-identical. The first mismatch
positions were
[217, 322, 293, 58]for both L2 and L3. Since the L2 and L3replays match each other exactly, Mooncake L3 does not introduce an additional
divergence. #35086 fixes the packed Mamba checkpoint index; the remaining
cold/replay divergence is the separate cold-prefill persistence-boundary issue
described in #33515.
Accuracy validation
Ran the repository's
benchmark/gsm8k/bench_sglang.pyagainst the same NEXTN +Mooncake service using the first 200 examples from the official GSM8K test set:
Result on the complete stack:
194/200correct (97.0%accuracy),0%invalidresponses, no request failures, and
218.519 slatency. The earlierincomplete-stack result was
195/200; only prompt 182 changed correctness,from the correct answer
23to24. This run includes #35086 but does notinclude #33515's cold-prefill alignment change.
CI States
Latest PR Test (Base): ⏳ Run #32013404865
Latest PR Test (Extra): ⏳ Run #32013404713