Skip to content

fix(npu): split packed MTP HiCache transfers by pool - #35013

Draft
momaekar1 wants to merge 22 commits into
sgl-project:mainfrom
momaekar1:fix/npu-hicache-mtp-packed-transfer
Draft

fix(npu): split packed MTP HiCache transfers by pool#35013
momaekar1 wants to merge 22 commits into
sgl-project:mainfrom
momaekar1:fix/npu-hicache-mtp-packed-transfer

Conversation

@momaekar1

@momaekar1 momaekar1 commented Aug 16, 2026

Copy link
Copy Markdown

Important

For standalone Ascend HiCache validation, this PR branch bundles the complete
commit history of #32500 through #32500@9acc6b733. Users can fork this branch
and 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-end
Qwen3.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_exchange path for page_first_direct, but that path invokes
the 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 + draft host-layer range, so the draft K/V layers are not transferred
through 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:

--tp-size 2
--attention-backend ascend
--device npu
--page-size 128
--enable-hierarchical-cache
--hicache-write-policy write_through
--hicache-storage-backend mooncake
--hicache-storage-prefetch-policy wait_complete
--speculative-algorithm NEXTN
--speculative-num-steps 3
--speculative-eagle-topk 1
--speculative-num-draft-tokens 4
  1. Generate from a 3000-token prompt to populate HiCache.
  2. For L2, generate a separate 6000-token filler to evict the original prefix
    from L1, then replay the original prompt.
  3. For L3, wait for write-through completion, call /flush_cache, then replay
    the original prompt.
  4. Request return_cached_tokens_details and inspect the target/draft transfer
    at the first kernel_ascend H2D 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

  • Skip CUDA-style packed device pointer arrays for the Ascend MTP path.
  • Iterate over the target pool and every draft pool on Ascend.
  • Slice the packed host K/V buffer to the layer range that belongs to each
    device pool.
  • Invoke transfer_kv_dim_exchange once per pool for both H2D and D2H.
  • Add focused tests for Ascend initialization and target/draft host-layer
    slicing.
  • Bundle all feat(hicache): support Ascend Mamba states with FIA and async IO #32500 commits through 9acc6b733 so this fork branch is directly
    runnable for Ascend HiCache validation.

Focused validation:

L2/L3 replay validation

Validation stack: latest main with #34560, #32500@9acc6b733, #35086, and
this fix. Model: Qwen3.6-27B BF16, TP2, Ascend NPU, NEXTN enabled, 4 prompts
with 3000 input tokens and 512 output tokens.

Replay tier Isolated cache details Result
Host L2 host=2944, device=0, storage=0 for 4/4 prompts passed
Mooncake L3 storage=2944, device=0, host=0 for 4/4 prompts passed

The L2 and L3 replay output IDs were identical for all four prompts:

sha256=a0349ca2a8cdf023da0ef7b86370fbc25eac55ba5d8a656665a845c4a1d969d9

Their speculative statistics were also identical:

  • weighted accept rate: 62.8169%
  • weighted accept length: 2.8845
  • correct/proposed drafts: 1338/2130

Cold 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 L3
replays 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.py against the same NEXTN +
Mooncake service using the first 200 examples from the official GSM8K test set:

num_questions=200
num_shots=5
parallel=4
temperature=0
max_new_tokens=512
chat_template_kwargs={"enable_thinking": false}

Result on the complete stack: 194/200 correct (97.0% accuracy), 0% invalid
responses, no request failures, and 218.519 s latency. The earlier
incomplete-stack result was 195/200; only prompt 182 changed correctness,
from the correct answer 23 to 24. This run includes #35086 but does not
include #33515's cold-prefill alignment change.


CI States

Latest PR Test (Base): ⏳ Run #32013404865
Latest PR Test (Extra): ⏳ Run #32013404713

@github-actions github-actions Bot added hicache Hierarchical Caching for SGLang npu labels Aug 16, 2026
@momaekar1
momaekar1 force-pushed the fix/npu-hicache-mtp-packed-transfer branch from 117edb4 to 98ae7de Compare August 17, 2026 02:49
@momaekar1
momaekar1 force-pushed the fix/npu-hicache-mtp-packed-transfer branch from 98ae7de to 06ad15e Compare August 17, 2026 09:04
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

hicache Hierarchical Caching for SGLang npu

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants