Qwen3-14B serving: add device top-k sampling

[zmnobug]

Summary

Extend the Qwen3-14B device sampling path introduced by #47 to support non-greedy top-k sampling.

For supported requests, a standalone NPU operator reduces full-vocabulary logits to a fixed-width candidate set. The existing host sampler consumes only the candidate values and token IDs, then applies temperature, top-p filtering, distribution validation, and multinomial sampling.

Closes #74.

Paired exact-selector pypto-lib PR: hw-native-sys/pypto-lib#787

What Changed

• Add an executor capability for the maximum device top-k width.
• Add SamplingCandidates to prefill/decode result contracts.
• Automatically enable device candidate selection when temperature > 0 and 0 < top_k <= 32.
• Preserve the existing CPU full-logits fallback for unsupported configurations.
• Compile the Qwen3 selector as a separate L3 callable.
• Reuse one selector for prefill greedy top-1 and non-greedy top-k candidate selection; decode greedy remains in the decode path introduced by Qwen3-14B serving: device-side embedding and sampling #47.
• Allocate reusable prefill/decode candidate buffers and return only active rows.
• Reuse the existing sampler for temperature, top-p, distribution validation, and multinomial selection.
• Wire both offline generation and the serving worker through the same capability-based path.
• Add routing, fallback, candidate sampling, row validation, and operator integration tests.

Scope

• Qwen3-14B on the NPU backend.
• Fixed device candidate capacity of 32; requests with a larger top_k use the CPU fallback.
• The selector remains a standalone operator and is not fused into prefill or decode.
• Final probabilistic sampling remains on the host.

Performance

Final exact top-k: latest main comparison

Measured on device 8 with /data/models/Qwen3-14B, prompt 北京故宫是, temperature=0.8, top_k=32, and top_p=1.0. The CPU baseline is current upstream/main@5d54073 with its pinned pypto-lib 402d12a. Both paths ran sequentially on the same device. Model initialization is excluded.

Twenty-token values are means over three runs:

Metric	CPU top-k baseline	Exact NPU top-k	Change
Generate E2E	1.571 s	0.749 s	-52.3%
E2E throughput	12.76 tok/s	26.72 tok/s	+109.4%, 2.09x
Prefill / TTFT	48.3 ms	51.3 ms	+3.0 ms
Decode API	34.8 ms/token	35.4 ms/token	+1.8%
Time outside prefill/decode APIs	861.3 ms	23.7 ms	-97.3%

The individual 20-token E2E runs were:

• CPU: 1.581, 1.668, and 1.465 s.
• Exact NPU: 0.742, 0.762, and 0.742 s.

A longer 128-token run confirmed that the E2E gain persists across 127 decode steps:

Metric	CPU top-k baseline	Exact NPU top-k	Change
Generate E2E	9.330 s	4.628 s	-50.4%
E2E throughput	13.72 tok/s	27.66 tok/s	+101.6%, 2.02x
Prefill / TTFT	50 ms	54 ms	+4 ms
Decode API	34.3 ms/token	35.5 ms/token	+3.5%
Time outside prefill/decode APIs	4.918 s	0.064 s	-98.7%

The standalone selector adds work inside run_decode, but it removes the dominant host full-vocabulary top-k and probability-processing work outside the executor APIs. This is why decode API time is approximately flat while generation E2E throughput roughly doubles versus CPU top-k.

Exactness Design

The follow-up exact implementation replaces the original per-512-chunk Top-4 approximation with a provably exact hierarchical reduction:

1. Split 151936 real-vocabulary logits into 74 full 2048-token groups and one 384-token tail group.
2. Compute an exact Top-32 for every group using sort32 and mrgsort.
3. Merge the resulting 75 x 32 = 2400 candidates in a 4096-entry padded buffer and select the exact global Top-32.
4. Compare the NPU output directly against torch.topk(logits[:, :REAL_VOCAB], 32).

The proof is straightforward: if an entry is not in its group's Top-32, at least 32 entries in the same group rank ahead of it, so it cannot belong to the global Top-32. Therefore the union of exact group Top-32 sets contains the exact global Top-32.

The test fixture covers both adversarial distributions:

• All 32 largest values concentrated in one 512-token region, which fails the original Top-4 approximation.
• One of the 32 largest values placed in each of 32 different chunks, which checks token-ID propagation.

The 2048-token grouping also reduces the generated sorting/merge task graph compared with computing Top-32 independently for all 297 512-token chunks. The exact implementation passed the shared-L3 batch-16 warmup and full prefill/decode generation without the 507018 timeout seen with larger task-graph prototypes.

Data Movement Note

The host sampler consumes only 32 values and 32 token IDs per active row instead of scanning the full vocabulary. However, prefill/decode result contracts still retain full logits for fallback, and the selector is a separate callable. Whether the runtime completely avoids a physical full-logits D2H/H2D round trip requires runtime copy tracing or a device-resident logits handle; this PR does not claim that transport has been eliminated.

Validation

• Exact Top-32 NPU output passed direct full-vocabulary torch.topk comparison.
• Greedy selection_k=1 regression passed host argmax comparison.
• tests/test_batching.py and tests/test_device_sampling_submission.py: 29 passed after rebase.
• git diff --check: passed in serving and pypto-lib.
• Shared-L3 batch-16 prefill/decode warmup: passed.
• Rebased standalone selector, 100 rounds: 428.0 us mean.
• Rebased Qwen3-14B 20-token exact top-k generation: 0.769 s E2E, 26.01 tok/s.
• Qwen3-14B 128-token exact top-k generation: passed.

[coderabbitai]

Important

Review skipped

Auto incremental reviews are disabled on this repository.

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: Organization UI

Review profile: CHILL

Plan: Pro Plus

Run ID: 6ec515c2-0f57-43a6-9b5b-70d04e5abb55

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

📝 Walkthrough

Walkthrough

Qwen3 device sampling now uses a shared top-k selector kernel, exposes candidate tensors through executor results, and performs final temperature/top-p sampling on the host. Capability checks preserve host fallback behavior for unsupported configurations.

Changes

Device top-k sampling

Layer / File(s)	Summary
Sampling contracts and fallback flow python/core/types.py, python/core/executor.py, python/core/engine.py, python/core/serving_worker.py, python/core/sampler.py	Adds SamplingCandidates, executor capability detection, batch flags, candidate-based sampling, and host fallback logic.
Qwen3 top-k kernel compilation examples/model/qwen3_14b/runner/npu_executor.py, examples/model/qwen3_14b/runner/qwen3_l3_dispatch.py	Replaces greedy-kernel compilation and wiring with topk_select, including fixed-shape validation and top-k buffers.
Qwen3 runner execution integration examples/model/qwen3_14b/runner/npu_runner.py	Runs selector-based greedy and top-k paths, returns candidates from prefill/decode, shares new buffers, and dispatches the top-k worker program.
Kernel and batching validation tests/test_batching.py, tests/test_device_sampling_submission.py, pypto-lib	Updates test doubles and assertions for candidate sampling, runtime vocabulary limits, tie-breaking, and top-k kernel wiring; advances the library submodule reference.

Estimated code review effort: 4 (Complex) | ~45 minutes

Sequence Diagram(s)

sequenceDiagram
  participant Request
  participant ServingWorker
  participant QwenRunner
  participant TopKKernel
  participant Sampler
  Request->>ServingWorker: submit top-k generation request
  ServingWorker->>QwenRunner: run prefill/decode with top-k enabled
  QwenRunner->>TopKKernel: dispatch topk_select
  TopKKernel-->>QwenRunner: return candidate values and token IDs
  QwenRunner-->>ServingWorker: return sampling_candidates
  ServingWorker->>Sampler: sample_from_candidates
  Sampler-->>ServingWorker: selected token

Loading

Poem

I’m a rabbit with kernels to hop,
Sorting bright candidates right to the top.
The sampler picks with a whisker-soft gleam,
While fallback paths guard the stream.
New buffers bloom in the NPU dream.

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Linked Issues check	✅ Passed	The changes add device top-k selection, candidate propagation, capability gating, and CPU fallback as requested by #74.
Out of Scope Changes check	✅ Passed	The modified files and tests are all scoped to Qwen3-14B device top-k sampling and its supporting runtime wiring.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Title check	✅ Passed	The title accurately summarizes the main change: adding device-side top-k sampling for Qwen3-14B serving.
Description check	✅ Passed	The description is directly related and matches the changeset's top-k device sampling, fallback, and serving flow updates.

---

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.}

[gemini-code-assist]

Code Review

This pull request introduces support for device-side top-k candidate selection and sampling for the Qwen3-14B model on NPU. It adds the necessary JIT compilation, host dispatching, and buffer allocations for the topk_select kernel. Additionally, it updates the core engine, serving worker, and sampler to support retrieving and sampling from these device-provided top-k candidates when configured. Unit tests are also added to verify the correctness of the top-k sampling path. Feedback suggests adding a boundary check for row_idx in sample_from_candidates to prevent potential out-of-bounds errors.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[coderabbitai]

🧹 Nitpick comments (1)

examples/model/qwen3_14b/runner/qwen3_l3_dispatch.py (1)

137-144: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Remove the stale qwen3_greedy_sample_host wrapper examples/model/qwen3_14b/runner/qwen3_l3_dispatch.py:137-144. The runtime now routes greedy prefill through _run_sampling_selector(..., selection_k=1) and only wires topk_select_fwd, so this shim still dereferences greedy_sample_fwd even though nothing assigns it. Update the batching test that uses this symbol as a slice marker if you drop it.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@examples/model/qwen3_14b/runner/qwen3_l3_dispatch.py` around lines 137 - 144,
Remove the unused qwen3_greedy_sample_host wrapper and any associated stale
greedy_sample_fwd reference from the dispatch module. Update the batching test’s
slice marker to use the next valid symbol or boundary so it remains correct
after the wrapper is deleted.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Nitpick comments:
In `@examples/model/qwen3_14b/runner/qwen3_l3_dispatch.py`:
- Around line 137-144: Remove the unused qwen3_greedy_sample_host wrapper and
any associated stale greedy_sample_fwd reference from the dispatch module.
Update the batching test’s slice marker to use the next valid symbol or boundary
so it remains correct after the wrapper is deleted.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro

Run ID: 5d08f065-eb24-4d7a-9795-94d920e769c3

📥 Commits

Reviewing files that changed from the base of the PR and between 49297c2 and f7d2775.

📒 Files selected for processing (11)

• examples/model/qwen3_14b/runner/npu_executor.py
• examples/model/qwen3_14b/runner/npu_runner.py
• examples/model/qwen3_14b/runner/qwen3_l3_dispatch.py
• pypto-lib
• python/core/engine.py
• python/core/executor.py
• python/core/sampler.py
• python/core/serving_worker.py
• python/core/types.py
• tests/test_batching.py
• tests/test_device_sampling_submission.py

[zmnobug]

Addressed the automated review feedback in the latest commits:\n\n- Added candidate row-count and row_idx bounds validation with negative/out-of-range tests in e23603f.\n- Removed the stale qwen3_greedy_sample_host wrapper and unused greedy_sample_fwd reference in cbf4e75.\n\nValidation: 28 focused tests passed and Ruff checks passed.

[zmnobug]

CI OOM follow-up: commit 6b3fe35 reduces the two Qwen3 CI guards from a 1 GiB to a 512 MiB PTO2_RING_HEAP. The runtime reserves roughly four heaps, so this releases about 2 GiB for the 40-layer model; the failed 7,130,316,800-byte allocation is the stacked BF16 FFN weight, not a sampling buffer.\n\nLocal validation with the new CI values completed static weight upload, prefill/decode warmup, KV allocation, and generation successfully (65.79 GB device: 33.64 GB peak non-KV, 2.28 GB arena). The new workflow run is waiting for maintainer approval: https://github.com/hw-native-sys/pypto-serving/actions/runs/29303581069

[zmnobug]

Local CI-equivalent validation passed with zm_pypto and the new 512 MiB ring heap configuration.\n\n- Test: tests/test_qwen3_accuracy.py\n- Result: 1 passed, 4 warnings in 94.29s\n- Device: 0\n- Task: task_20260713_221343_20357588630\n\nThe earlier local tbe error was caused by overwriting CANN PYTHONPATH; preserving it as PYTHONPATH=/data/zhaomin/pypto-serving:$PYTHONPATH resolves the environment issue.