[eeeeeio]

Continuation of #38197 (closed automatically when its branch was rebuilt
during the rebase onto the 10/21 kernels/ops refactor; no review had started
there). Same three-file change, rebased onto current main.

Motivation

NGRAM speculative decoding is currently restricted to a platform allowlist
(_handle_ngram rejects non-CUDA/CPU/XPU devices). The restriction is an
allowlist check rather than a fundamental dependency: with three
platform-scoped changes, the full NGRAM path (draft lookup → chain build →
target verify → accept-path KV relocation) runs end-to-end on Ascend 910C
(task: Ascend#1109).

This extends speculative-decoding coverage on NPUs, complementing the DFlash2
NPU enablement (#35629, merged) and the NGRAM roadmap (#21052). It is opt-in
and changes nothing for CUDA/CPU/XPU users.

Modifications

1. srt/arg_groups/speculative_hook.py — allow "npu" in the
   _handle_ngram device allowlist.

2. kernels/ops/speculative/reconstruct_tree.py — the NPU path has no
   device kernel for reconstruct_indices_from_tree_mask (sgl_kernel ships
   CUDA-only binaries, and the module-level import would fail on NPU), so this
   adds a host fallback and routes it in reconstruct_indices_from_tree_mask_func
   next to the existing XPU Triton branch, following the same dispatch pattern
   as the HIP routing in [ROCm][Spec] Route NGRAM tree reconstruction to the JIT kernel on HIP #43328 (device kernel there, host fallback here). The fallback reproduces the kernel contract into
   the caller's preallocated device tensors (slot identity, tree depth +
   seq_len positions, and the tree links); CUDA/CPU/XPU dispatch is unchanged.

3. srt/hardware_backend/npu/memory_pool_npu.py — the speculative accept
   path calls move_kv_cache(tgt_loc, src_loc), but the base tiled-copy
   implementation requires _kv_copy_config, which the NPU paged buffer layout
   never builds (see _init_kv_copy_and_warmup). Add an NPU-native
   move_kv_cache that relocates slots with plain indexing, viewing each layer
   as flat [num_slots, head_num, head_dim] (V uses v_head_dim) — the same
   convention get_cpu_copy/load_cpu_copy already use for both buffer
   layouts.

Accuracy Tests

• Functional verification on hardware (see environment below): Qwen3-0.6B
  bf16 served with NGRAM enabled returns coherent greedy chat completions,
  with the server stable across repeated requests and the NPU-native
  move_kv_cache invoked on every accept step.
• The host fallback reproduces the device kernel's documented contract
  (same four outputs, in-place mutation), mirroring the host tree-link
  derivation used by the grammar path (kept as a local copy to avoid a
  circular import).
• Lossless by construction in greedy decoding: draft tokens are verified by
  the target model, so accepted tokens come from the target distribution and
  output is expected to be token-identical to non-speculative greedy
  decoding. A dedicated accuracy benchmark (e.g. GSM8K) was not run in this
  setup.

Speed Tests and Profiling

• Benchmark (sglang.bench_serving, random 128/64, 32 prompts, concurrency 8,
  request-rate 4; 3 runs per config, the first on-run serving as warm-up):
  decode latency (3-run Median TPOT) drops from ~28.0 ms to ~13.7 ms (−51%)
  with NGRAM on.
• Output throughput is unchanged within run-to-run variance on this
  host-bound 0.6B setup (per-run values ranged ~119–141 tok/s across on/off
  runs). TTFT +21–30 ms across paired runs (per-request n-gram state init).
  The host-fallback overhead is included in all reported numbers.

Known limitations (pre-existing, not introduced by this PR)

• Chain verify only for now: tree verify (max_bfs_breadth > 1) requires
  page_size == 1, which the FIA kernel rejects (blockSize must be a
  multiple of 16 in [16, 1024]; see [Bug] [NPU] Ascend FIA attention should reject page_size=1 or document block_size alignment requirement #25169). With
  --speculative-ngram-max-bfs-breadth 1 the feature works end-to-end;
  lifting the kernel constraint is follow-up work.
• Use --page-size 128: page_size 16 silently corrupts output on the
  ascend backend — pre-existing, tracked in [Bug] [NPU] ascend attention backend silently corrupts output with page_size 16 (128 OK, 1 rejected per #25169) #38196 (independent of
  speculative decoding); --page-size 128 is the correct setting, as on main.
• A separate vendor-kernel crash in
  sgl_kernel_npu.norm.split_qkv_rmsnorm_rope (hit by any Qwen3 NPU run on
  this stack) was worked around locally with the non-fused path; unrelated to
  NGRAM, and will be filed as a separate upstream issue.

Environment (hardware verification, 2026-09)

• Hardware: 2× Ascend 910C (Atlas 800 A3, dual-die; single-card run), CANN 9.0.0
• Note: this PR was rebased onto main (2026-10-22) to follow the
  kernels/ops speculative-dispatch refactor ([Refactor] spec: Move platform dispatch and Triton kernels into kernels/ops #45504). The host-fallback
  logic is unchanged from the originally hardware-verified implementation
  (module moved, logic identical); the numbers above were measured with that
  logic, and no hardware re-run was done after the rebase.
• Stack: python 3.11, torch 2.10.0 + torch_npu 2.10.0, sgl-kernel-npu 2026.6.1,
  triton-ascend 3.2.1.dev; sglang source build
• Reproduce: --attention-backend ascend --disable-overlap-schedule --disable-cuda-graph --speculative-algorithm NGRAM --speculative-num-draft-tokens 4 --speculative-ngram-max-bfs-breadth 1 --page-size 128 on any Qwen3 NPU deployment.

---

CI States

Latest PR Test (Base): ❌ Run #40914727809
Latest PR Test (Extra): ❌ Run #40914727795
Latest PR Test (AMD ROCm 10): ❌ Run #40914727849

[eeeeeio]

Fork PR — CI jobs were skipped by default. Could a maintainer please help
trigger CI (e.g. /tag-run-ci-label), especially the NPU-relevant checks
such as pr-test-npu-finish? (Other *-finish gates showing red are just
the fork-PR skip cascade, not real failures.)

[gjsheu]

Your PR does not meet the requirements of the issue. Please verify it promptly.

[eeeeeio]

Hi @gjsheu, thanks for the review. To clarify: the skill deliverables required by this issue (SKILL.md, prompt/invocation templates, baseline-collection scripts, report generation, and recommended launch commands) are submitted to a dedicated skills repository as guided in the task discussion group — see LinyuanLi0046/Ascend-SGLang-Skills#23

This PR covers the code-level part: it upstreams the NGRAM switch from this task's feature evaluation (speculative decoding on Ascend NPU). A full status update covering both PRs is posted on the issue: Ascend#1109 (comment)

Happy to adjust if anything specific is missing.