Optimize Qwen3 device sampling outputs

[zmnobug]

Summary

• keep Qwen3 device-greedy logits outputs worker-resident in the serving runner
• avoid copying full prefill/decode logits back to host when sampled_token_ids is already produced on device
• keep decode next_hidden on the original compiled host/shared buffer path to minimize the behavioral diff
• return an empty CPU logits placeholder for greedy device-sampling paths
• keep result sampling decisions in one helper by passing result objects instead of separate logits / sampled_token_ids arguments
• add behavioral regression coverage for worker-resident output caching and logits placeholder handling

Test

/data/zhaomin/.conda/envs/zm_pypto/bin/python -m pytest tests/test_device_sampling_submission.py
/data/zhaomin/.conda/envs/zm_pypto/bin/python -m pytest tests/test_batching.py

Result:

tests/test_device_sampling_submission.py: 5 passed
tests/test_batching.py: 18 passed

Qwen3 serving guard was also run locally with zm_pypto:

tests/test_qwen3_serving.py: 4 passed
tests/test_qwen3_serving.py::test_prefix_cache_reuses_prefix_and_preserves_output: 1 passed

Benchmark

Qwen3-14B, full 40-layer, model path /data/models/Qwen3-14B, A2A3, prompt 北京故宫是, 128 generated tokens, offline npu_generate.py --profile.

metric	baseline	this PR	delta
generate e2e	5.37s	4.99s	-7.08%
throughput	23.83 tok/s	25.64 tok/s	+7.60%
run_decode total	5.07s	4.67s	-7.89%
kernel.decode_layer avg	36.61 ms	31.36 ms	-14.34%

Key raw timing:

baseline: generate=5.37s, decode=5.07s, decode kernel avg=36.61ms, throughput=23.83 tok/s
this PR:  generate=4.99s, decode=4.67s, decode kernel avg=31.36ms, throughput=25.64 tok/s

[coderabbitai]

📝 Walkthrough

Walkthrough

This PR modifies the Qwen3-14B NPU runner to keep large decode/prefill output tensors (logits, next-hidden) resident on the worker device instead of transferring to host, gated by an environment variable and allow_device_greedy_sampling flag. Adds a caching layer for these buffers and a validating unit test.

Changes

Worker-resident output tensors

Layer / File(s)	Summary
Data shapes and gating setup examples/model/qwen3_14b/runner/npu_runner.py	Adds os import, widens _DecodeKernelInputs.logits/next_hidden to accept torch.Tensor | DeviceTensor, and adds an _l3_output_tensors cache field.
Gating and buffer allocation helpers examples/model/qwen3_14b/runner/npu_runner.py	Adds _use_worker_resident_outputs, _output_kernel_arg, and _result_logits helpers to gate, allocate/cache, and convert worker-resident output buffers.
Prefill/decode integration examples/model/qwen3_14b/runner/npu_runner.py	Updates run_prefill, run_decode, _pad_decode_inputs, _maybe_run_sample_embed, and _prefill_kernel_args to select and pass device-resident logits/next-hidden buffers.
Cache cleanup and test coverage examples/model/qwen3_14b/runner/npu_runner.py, tests/test_device_sampling_submission.py	Clears the output tensor cache in close() and adds a unit test validating worker-resident output handling markers exist in the runner source.

Estimated code review effort: 3 (Moderate) | ~25 minutes

Possibly related PRs

• hw-native-sys/pypto-serving#29: Both modify npu_runner.py's run_prefill/run_decode worker-resident tensor handling and caching.
• hw-native-sys/pypto-serving#47: Both extend the same runner to carry sampling outputs (logits/next-hidden) through the device path instead of host sampling.

Poem

A rabbit hops through logits deep,
No more to host they need to leap,
Cached on worker, buffers stay,
Greedy tokens find their way,
Thump-thump, the tensors home to keep! 🐇✨

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Title check	✅ Passed	The title clearly summarizes the main change: optimizing Qwen3 device sampling outputs.
Description check	✅ Passed	The description is directly aligned with the worker-resident output changes, tests, and benchmarks in the PR.

---

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 optimizes performance during greedy generation by keeping large output tensors (such as logits and next hidden states) resident on the NPU worker, thereby avoiding unnecessary host-device transfers. A test case has also been added to verify this behavior. The feedback identifies a potential device mismatch issue in the _result_logits helper, where returning an empty placeholder on the NPU device instead of the CPU could lead to downstream runtime errors. It is recommended to instantiate this placeholder directly on the CPU.

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)

tests/test_device_sampling_submission.py (1)

71-79: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Test only greps source text; doesn't exercise runtime behavior.

This test asserts literal substrings exist in the runner's source file rather than exercising _use_worker_resident_outputs, _output_kernel_arg caching, or _result_logits's actual output for both DeviceTensor and host-tensor inputs. It will pass even if the gating logic is broken (e.g., wrong boolean condition, env var check inverted) as long as the strings remain present, and will break on harmless refactors (renames, reformatting) that don't change behavior. Since the PR description cites this test as the primary regression coverage for this optimization, consider adding a behavioral unit test that mocks the worker (similar to the WorkerTensor/alloc_tensor pattern already used in tests/test_batching.py) to verify: _use_worker_resident_outputs respects QWEN14B_DEVICE_OUTPUTS=0/1/unset, _output_kernel_arg returns a cached instance on repeat calls with the same key, and _result_logits returns the correct empty vs. sliced tensor for each input type.

🤖 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 `@tests/test_device_sampling_submission.py` around lines 71 - 79, The current
test only checks for source substrings and does not verify the actual
device-output behavior. Replace or supplement
test_device_greedy_keeps_large_outputs_worker_resident with a behavioral test
that exercises _use_worker_resident_outputs, _output_kernel_arg, and
_result_logits on npu_runner.py using a mocked worker/tensor setup like the
existing WorkerTensor/alloc_tensor patterns in tests/test_batching.py. Cover
QWEN14B_DEVICE_OUTPUTS unset/0/1, confirm _output_kernel_arg caches per key, and
assert _result_logits returns the expected empty or sliced tensor for both
DeviceTensor and host-tensor inputs.

🤖 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 `@tests/test_device_sampling_submission.py`:
- Around line 71-79: The current test only checks for source substrings and does
not verify the actual device-output behavior. Replace or supplement
test_device_greedy_keeps_large_outputs_worker_resident with a behavioral test
that exercises _use_worker_resident_outputs, _output_kernel_arg, and
_result_logits on npu_runner.py using a mocked worker/tensor setup like the
existing WorkerTensor/alloc_tensor patterns in tests/test_batching.py. Cover
QWEN14B_DEVICE_OUTPUTS unset/0/1, confirm _output_kernel_arg caches per key, and
assert _result_logits returns the expected empty or sliced tensor for both
DeviceTensor and host-tensor inputs.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro

Run ID: 0bb22df9-dda9-48c3-8aa4-fd82de3fc8a1

📥 Commits

Reviewing files that changed from the base of the PR and between 5325cf2 and 6a3e55b.

📒 Files selected for processing (2)

• examples/model/qwen3_14b/runner/npu_runner.py
• tests/test_device_sampling_submission.py

[zmnobug]

Addressed the review feedback in 83eec9b:

• _result_logits now creates the empty greedy placeholder on CPU.
• Replaced the source-grep check with a behavioral test covering QWEN14B_DEVICE_OUTPUTS, output tensor caching, and host/DeviceTensor _result_logits behavior.

Validation:

/data/zhaomin/.conda/envs/zm_pypto/bin/python -m pytest tests/test_device_sampling_submission.py
5 passed