Feat/shadownpu sparse attention

[yinwangsong]

This is an implementation of MobiSys 2026 paper ShadowNPU: System and Algorithm Co-design for NPU-Centric
On-Device LLM Inference, which is coauthored by PKU and BUPT team.

Summary

This PR adds a reproducible ShadowNPU-style sparse prefill attention path for
Qwen models on Qualcomm Hexagon NPU devices. It combines static HMX INT8 QK,
CPU INT8 Top-K selection, sparse QK/softmax/PV, and pipelined CPU/NPU
execution while keeping decode attention dense.

The PR packages one-command reproduction for Qwen2.5-1.5B and provides the
source compatibility paths required by Qwen1.5-1.8B and Qwen1.5-0.5B.

Motivation

Dense prefill attention becomes increasingly expensive as the prompt length
grows. This change introduces a static, model-calibrated sparse attention path
that reduces prefill attention work while preserving strict retrieval quality
for the supported Qwen2.5-1.5B and Qwen1.5-1.8B configurations.

What changed

Runtime and scheduling

• Add deterministic per-head retention scheduling.
• Add HMX INT8 QK selection with static Q/K and target-requant scales.
• Add CPU INT8 Top-K and sparse QK/softmax/PV implementations.
• Add FIFO and profile-driven pipeline scheduling across NPU, Top-K, packing,
  and sparse attention workers.
• Add attention timing and pipeline profiling support.

QNN and model compatibility

• Integrate the sparse prefill path with the QNN backend.
• Support Qwen2.5-1.5B M256 x 2 execution for a logical M512 attention chunk.
• Add Qwen1.5 activation-scale overrides and final-token CPU refinement.
• Preserve dense decode attention and model-specific QNN op-package handling.

Tooling and reproducibility

• Add HMX operator-bank build and validation utilities.
• Add head-sparsity profiling, scale calibration, and pipeline-latency tools.
• Add an Android benchmark runner with alternating dense/sparse order,
  strict retrieval checks, and median aggregation.
• Add focused C++ and Python tests and end-to-end reproduction documentation
  in docs/shadownpu_sparse_attention.md.

Supported configurations

Model	QNN chunking	HMX bank	Reproduction scope
Qwen2.5-1.5B	M256 x 2	H12, M512/K128/N4160 target-requant	Packaged runner
Qwen1.5-1.8B	M256	H16, M256/K128/N4160 target-requant	External model/context/profile required
Qwen1.5-0.5B	M128	H16, M128/K64/N2048 or N4160 target-requant	External model/context/profile required

Device validation

• Device: Xiaomi 14 (houji, SM8650)
• OS: Android 16 / API 36
• Canonical PR commit: 7e16701819ed3540c1d6aa35a8bec0c731cb08de
• Measurement origin before squash: 3b384c34525e68c423e336af50401df239101dbf
• Shared source tree: 7d739aad6cec09366c09fdd93afe71dbf766ce04
• Benchmark binary SHA-256: a8ee27d0292db9d70fe0dd4fe224124d8a24fa41d492a16cc8f649bd5aa3ffbb
• Measurements: 54 formal runs, 3 repetitions per model/prompt/mode

The measurement commit and the canonical PR commit have an identical Git
tree. The original measurement commit is retained above for provenance
because the branch was squashed after the device runs.

Each formal log contains exactly one BENCH_RESULT, QUALITY_RESULT, and
QWEN_ATTN_PROFILE, observes the FastRPC redirect, and contains no QNN or
FastRPC error marker. Attention mode, chunking, actual retention, CPU
placement, model compatibility settings, and process exit status were also
validated.

Performance

TTFT values are medians of three runs. Speedup is dense TTFT divided by sparse
TTFT within the same row.

Model	Prompt tokens	Dense TTFT	Sparse TTFT	Speedup	Actual retention	Quality
Qwen2.5-1.5B	1024	5320.753 ms	2514.360 ms	2.116x	0.139438	PASS (3/3 each mode)
Qwen2.5-1.5B	2048	16391.518 ms	6829.007 ms	2.400x	0.138950	PASS (3/3 each mode)
Qwen2.5-1.5B	4096	102762.251 ms	27046.845 ms	3.799x	0.138706	PASS (3/3 each mode)
Qwen1.5-1.8B	1024	6858.861 ms	2961.117 ms	2.316x	0.166194	PASS (3/3 each mode)
Qwen1.5-1.8B	2048	20960.720 ms	7394.772 ms	2.835x	0.165706	PASS (3/3 each mode)
Qwen1.5-1.8B	4096	77908.256 ms	27037.667 ms	2.881x	0.165462	PASS (3/3 each mode)
Qwen1.5-0.5B	1024	4108.120 ms	1607.499 ms	2.556x	0.166193	PASS (3/3 each mode)
Qwen1.5-0.5B	2048	14647.205 ms	4948.297 ms	2.960x	0.165706	Diagnostic failure (0/3 each mode)
Qwen1.5-0.5B	4096	57411.790 ms	18436.579 ms	3.114x	0.165461	Diagnostic failure (0/3 each mode)

The Qwen1.5-0.5B 2048- and 4096-token rows are retained as diagnostic timing
only because strict retrieval fails in both dense and sparse modes. Absolute
latency is temperature-sensitive. Historical preserved rows in the document
also include legacy fixed-output banks, whereas the current reproduction uses
the recommended target-requant banks.

Reproduction

Build the Android/QNN benchmark and follow
docs/shadownpu_sparse_attention.md for artifact provisioning. The packaged
Qwen2.5 A/B runner is:

ADB_SERIAL=<serial> \
DEVICE_ROOT=/data/local/tmp/mllm \
RUN_DIR=/data/local/tmp/mllm/qwen-run \
HMX_BANK=/data/local/tmp/mllm/hmx-qwen2-m512-h12-n4160 \
REPETITIONS=3 \
tools/shadownpu/run_qwen_ab.sh

Test plan

• Build benchmark_qwen_npu with the Android/QNN configuration.
• Verify matching host/device benchmark SHA-256.
• Run 54 dense/sparse device measurements on Xiaomi 14.
• Validate strict retrieval for Qwen2.5-1.5B and Qwen1.5-1.8B.
• Preserve and label Qwen1.5-0.5B long-context diagnostic failures.
• Run the full host MLLM_TEST suite on the squashed commit.
• Run the focused ShadowNPU C++ tests on the squashed commit.
• Run all three Python unit-test suites on the squashed commit.
• Pass upstream GitHub Actions.

Known limitations

• Qwen1.5 reproduction requires separately provisioned model, context,
  profile, and operator-bank artifacts.
• Performance results are device-, temperature-, and QNN-version-dependent.
• The standard GitHub-hosted CI cannot execute Qualcomm HMX/FastRPC device
  benchmarks; those checks require a manually triggered self-hosted runner or
  maintainer device reproduction.

[coderabbitai]

Important

Review skipped

Auto reviews are disabled on base/target branches other than the default branch.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro Plus

Run ID: 78a8ddc0-fa66-459c-a1dc-c7104c950982

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}