feat(qwen3): use device embedding for prefill

[lyfne123]

Summary

• pass packed token IDs and the shared device-resident embedding table to Qwen prefill
• replace the prefill BF16 hidden buffer with an int32 token ID buffer
• skip host prompt embedding lookup for executors that support device embedding
• preserve host embedding input for other model runners
• pin pypto-lib to merged #774 commit 402d12a61dfd3f415ca9ec5356f9e7ff876b6ad8

Validation

• python -m pytest tests/test_batching.py tests/test_parallel.py -q (28 passed)
• ruff, header, and English-only checks passed
• full Qwen3-14B NPU compile and generation succeeded (4096-token warmup plus 5 generated tokens, exit 0)

[gemini-code-assist]

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

[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

Run ID: 72554f7b-300e-424b-a8c0-608f43d43b6a

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 prefill execution now passes token IDs and embedding weights to device kernels, allowing host embeddings to be omitted. Core batching types, engine and worker paths, compiled buffers, dispatch signatures, and tests were updated accordingly.

Changes

Device-side embedding prefill

Layer / File(s)	Summary
Optional prefill embeddings python/core/engine.py, python/core/serving_worker.py, python/core/types.py, python/core/executor.py	Embedding tensors are conditionally omitted for device-embedding executors, and PrefillBatch.input_embeddings is nullable.
Qwen3 token-ID prefill buffers examples/model/qwen3_14b/runner/npu_runner.py, examples/model/qwen3_14b/runner/npu_executor.py, pypto-lib	Qwen3 prefill workspaces, preparation, warmup, sharing, and arguments now use int32 token IDs instead of hidden states.
Prefill dispatch contract examples/model/qwen3_14b/runner/qwen3_l3_dispatch.py, examples/model/qwen3_14b/runner/npu_executor.py	The host wrapper and compilation arguments forward input IDs and embedding weights under the revised prefill signature.
Behavioral validation and compatibility tests/test_batching.py, examples/model/deepseek_v4/runner/npu_runner.py	Tests validate token-ID packing, skipped host lookups, dispatch inlining, and required embeddings for DeepSeek prefill.

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

Sequence Diagram(s)

sequenceDiagram
  participant Engine
  participant Worker
  participant Qwen3PrefillHost
  participant DeviceKernel
  Engine->>Worker: create PrefillBatch with token IDs
  Worker->>Qwen3PrefillHost: pass token IDs and optional embeddings
  Qwen3PrefillHost->>DeviceKernel: forward input IDs and embed weight
  DeviceKernel->>DeviceKernel: embed token IDs and execute prefill

Loading

Poem

A rabbit hops through tokens bright,
With kernels weaving IDs just right.
No host-side embeddings in the queue,
The device knows what it should do.
Tests thump softly: “All paths true!”

🚥 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 states the main change: using device embeddings for Qwen3 prefill.
Description check	✅ Passed	The description matches the changeset and summarizes the device-embedding refactor and pypto-lib pin.

---

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

[superxf]

supports_device_embedding还有必要么，可以做成一个通用流程？