Add native Qwen3 14B A8W8 serving path - #48
Conversation
📝 WalkthroughWalkthroughAdds an A8W8/BF16 checkpoint loader and decode-backend CLI options for Qwen3-14B, introduces L3-mode and A8W8-decode-loop execution paths with quantization-aware compilation and weight handling in the NPU executor, threads a ChangesA8W8/L3 Decode Feature
Dependency-free Qwen Tokenizer Fallback
Sequence Diagram(s)sequenceDiagram
participant Engine
participant Executor as Qwen314BPyptoExecutor
participant A8W8Loop as run_generate_a8w8_decode_loop
participant L3 as run_generate_l3
Engine->>Executor: try_generate_batch(record, requests, prefill_batch, config)
Executor->>Executor: validate_generate_batch / prompt_allocation_length
alt a8w8_decode_loop enabled
Executor->>A8W8Loop: run_generate_a8w8_decode_loop(...)
else l3_mode enabled
Executor->>L3: run_generate_l3(...)
end
Executor-->>Engine: GenerateResult(finish_reason)
sequenceDiagram
participant main as npu_generate.main
participant Loader as Qwen3A8W8DirectoryLoader
participant Index as _SafeTensorIndex
participant Runtime as RuntimeModel
main->>Loader: load(request)
Loader->>Loader: read config.json, build tokenizer
Loader->>Index: scan model-*.safetensors shards
Index-->>Loader: tensor offsets/dtype/shape
Loader->>Index: load(key) per tensor
Loader->>Loader: convert layouts (int8 kernel, bf16 dequant)
Loader->>Runtime: populate embeddings/norm/head/layers
Loader-->>main: LoadedModel
Estimated code review effort🎯 4 (Complex) | ⏱️ ~75 minutes Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
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. Comment |
There was a problem hiding this comment.
Code Review
This pull request introduces support for qwen3-a8w8 model quantization, including a new loader for compressed-tensors checkpoints and updated NPU execution and runner logic to handle A8W8 kernels and KV cache scaling. It also adds a dependency-free byte-level BPE tokenizer fallback. The review highlights significant code duplication in the NPU executor and runner logic, suggesting refactoring to improve maintainability. Additionally, it recommends performance optimizations for file I/O and tokenization, replacing SimpleNamespace with dataclass for better type safety, and avoiding bare except blocks.
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.
b09f2cf to
2343c8e
Compare
There was a problem hiding this comment.
Actionable comments posted: 4
🤖 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.
Inline comments:
In `@examples/model/qwen3_14b/npu_generate.py`:
- Around line 421-422: The QWEN_A8W8_BF16_KV environment flag is only ever set
and never cleared in main(), so a previous bf16-kv run can leak into later A8W8
runs in the same process. Update the argument-handling path around
args.model_format and args.decode_backend to deterministically reset
QWEN_A8W8_BF16_KV for every invocation, explicitly setting it to "1" only for
qwen3-a8w8 with bf16-kv and removing or clearing it otherwise.
In `@examples/model/qwen3_14b/runner/a8w8_loader.py`:
- Around line 86-89: Reshape the linear scale tensor in _hf_linear_to_bf16
before multiplying so it broadcasts over output channels correctly instead of
the last axis. Update the _hf_linear_to_bf16 flow to explicitly reshape or
unsqueeze the loaded scale from index.load(scale_key) before the weight * scale
operation, then keep the transpose/contiguous BF16 conversion as-is.
In `@python/core/tokenizer.py`:
- Around line 248-250: The byte-level tokenizer input is being altered in
Tokenizer.encode by NFC normalization, which breaks round-tripping for
decomposed Unicode text. Remove the unicodedata.normalize("NFC", text) step from
encode and preserve the original input bytes before byte-level BPE processing;
keep the rest of Tokenizer.encode and decode behavior unchanged.
- Around line 135-181: Add parity tests for the Qwen fallback tokenizer behavior
in _split_qwen_text, focusing on edge cases where the current split logic can
diverge from Hugging Face: repeated spaces, newline handling, and
punctuation/space boundaries. Create coverage that compares the fallback output
against the expected Qwen regex-based tokenization behavior for representative
inputs, so regressions are caught before relying on this path in serving.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 5243ea88-e297-4ca7-9abc-c6e9f3bcda03
📒 Files selected for processing (8)
examples/model/qwen3_14b/npu_generate.pyexamples/model/qwen3_14b/runner/a8w8_loader.pyexamples/model/qwen3_14b/runner/npu_executor.pyexamples/model/qwen3_14b/runner/npu_runner.pypython/core/pypto_executor.pypython/core/tokenizer.pypython/runtime/worker.pytests/test_batching.py
06d79c8 to
dd81f52
Compare
|
Follow-up on the two addressed serving review comments:
Validation: 48-token Qwen3 A8W8 generation remained normal, TTFT 7.054s, TPOT 67.8 ms/token. |
77e943b to
e6a73fd
Compare
e6a73fd to
9a34ac6
Compare
## Summary Add the Qwen3-14B A8W8 kernel implementation used by the native serving path. This PR keeps the existing BF16 Qwen3 path isolated and introduces separate A8W8 prefill/decode modules for the quantized model: - add `prefill_hidden` support for Qwen3-14B A8W8 hidden-state prefill chunks - add the optimized A8W8 `decode_fwd` layer kernel used by serving decode - handle INT8 weights, activation/weight scales, paged KV cache scale metadata, and A8W8-specific decode constants inside the A8W8 modules - keep scheduling, model loading, tokenizer, and request orchestration in `pypto-serving` rather than carrying standalone lib-side runners - keep the ordinary BF16 Qwen3-14B execution path separate from the A8W8 path ## Implementation Notes The lib-side deliverable is intentionally limited to kernels and kernel-adjacent compatibility: - `models/qwen3/14b/prefill_fwd_a8w8.py` implements the A8W8 prefill hidden path - `models/qwen3/14b/decode_layer_a8w8.py` implements the A8W8 decode layer path used by serving - `golden/runner.py` compatibility changes keep generated/runtime artifacts runnable for the JIT path - obsolete standalone debug/golden entry points were removed during slimming; debug-stage switches in the kernel remain available for targeted numerical diagnosis ## Validation End-to-end validation was run through the native serving stack with this kernel PR plus the matching serving/backend PRs: - prompt: `介绍一下北京故宫` - generated tokens: 48 - output quality: normal Chinese continuation - TTFT: `7.054s` - TPOT: `67.8 ms/token` - decode throughput: `14.76 tok/s` Focused serving and PyPTO checks also passed in the matching PRs. ## Related PRs / Issues - Tracking issue: #665 - Serving-side PR: hw-native-sys/pypto-serving#48 - PyPTO backend/lowering PR: hw-native-sys/pypto#1920 Co-authored-by: vegetabledoww <vegetabledoww@users.noreply.github.com>
39924ca to
2a24e51
Compare
|
Maintainer note: Before merging this PR, please update the |
8a2b15f to
91aae68
Compare
51a55d4 to
8c4d493
Compare
d2141ea to
b87e3d3
Compare
Add the compressed-tensors A8W8 loader for Qwen3-14B checkpoints and extend the common PyPTO runtime plumbing for PTO ISA pinning.
Add the Qwen3-14B A8W8 L2 runner responsible for paged KV cache management, child-memory argument preparation, prefill dispatch, decode dispatch, and logits collection.
Add the A8W8 PyPTO executor, connect npu_generate.py to the qwen3-a8w8 model format, keep BF16 and A8W8 runtime settings isolated, and add regression coverage for loader, runner, and CLI behavior.
Route Qwen3-14B A8W8 prefill and decode through HOST-level L3 wrappers, preallocate shared IO buffers for the L3 runner, and remove the A8W8 L2 callable fallback from the serving path. Consolidate A8W8 HOST wrappers into the existing Qwen3 L3 dispatch module. Use a closure-based factory to keep A8W8 kernel dependencies isolated from the BF16 dispatch globals. Validation: Python compile and pre-commit checks passed; 8-token A8W8 generation produced expected text at 55.1 ms/token, and prior 48-token L3-only validation succeeded.
Use BF16 KV caches and remove the obsolete per-row KV scale buffers to match the updated pypto-lib kernel ABI. Keep Gate and Up weights in INT8 form, propagate their scales through loading, compilation, and L3 dispatch, and update placeholder tensors accordingly. Refresh batching assertions for the split prefill embedding paths. Validated with: pytest -q tests/test_batching.py (37 tests completed; the local NPU runtime hung during process teardown).
Remove the obsolete pto_isa_commit CLI and executor plumbing now that PyPTO resolves managed PTO-ISA revisions from its runtime pin. Validation: pre-commit passed; 12 focused batching tests passed; deterministic 48-token A8W8 generation preserved the expected token IDs at 42.7 ms/token.
b87e3d3 to
5d8e271
Compare
Add explicit Qwen3 A8W8 model-format and tokenizer configuration to the online serving CLI, and propagate both settings through EngineConfig to worker processes. Select the compressed-tensors A8W8 loader and executor in the serving worker. Preserve the physical runtime batch size of 16 while allowing the online request concurrency to be configured independently. Validate the supported single-device topology, page size, runtime and KV dtypes, sequence limits, and checkpoint quantization method at startup so incompatible configurations fail before model execution. Validation: - Ruff, header, English-only, and diff checks pass - 64 relevant unit tests pass, including Qwen3, HTTP, BF16, and DeepSeek coverage - Device 0 completion, chat, and two-request concurrency tests pass - 48-token output matches the offline golden result

Summary
Add a native Qwen3-14B A8W8 serving path that runs compressed W8A8 checkpoints through both the standard
LLMEngineworkflow and the OpenAI-compatible HTTP service.The implementation covers the complete serving pipeline—from checkpoint loading and kernel compilation to L3 execution, paged KV-cache management, offline generation, and online completion/chat requests—while keeping the existing BF16 Qwen3 path isolated and unchanged.
Related to #52.
Uses the Qwen3 A8W8 kernel ABI landed by the merged pypto-lib PR #810.
What changed
Model format and checkpoint loading
qwen3-a8w8as a serving model format.A8W8 executor and kernel compilation
Qwen314BA8W8PyptoExecutor.pypto-libsource.L3 model runner and dispatch
DistributedProgramHOST wrappers.Offline generation integration
--model-format qwen3-a8w8inexamples/model/qwen3_14b/npu_generate.py.Online OpenAI-compatible serving
--model-format auto|qwen3-14b|qwen3-a8w8in the main serving CLI.--tokenizerso an A8W8 checkpoint can use a separate compatible Qwen3 tokenizer and official chat template.EngineConfiginto worker processes.--max-num-seqs./health,/v1/models,/v1/completions, and/v1/chat/completions.Commit breakdown
The PR remains intentionally split into seven reviewable commits:
feat(qwen3-a8w8): add loader and L2 runtime supportfeat(qwen3-a8w8): add A8W8 L2 model runnerfeat(qwen3-a8w8): compile kernels and wire generation CLIqwen3-a8w8tonpu_generate.pyand adds regression coverage.Use L3 dispatch for Qwen3 A8W8DistributedProgramdispatch.feat(qwen3-a8w8): align serving with optimized kernelspypto-libpin and associated runner, placeholder, dispatch, and test contracts.fix(qwen3): adapt serving to latest PyPTO runtimefeat(qwen3): support A8W8 online servingCompatibility and boundaries
pypto-libsubmodule is pinned to7a63278, which provides BF16 KV-cache interfaces, native INT8 Gate/Up inputs and scales, and explicit RNE narrowing.--tokenizer /data/models/Qwen3-14B.Performance
Offline decode
Recent deterministic single-device A2/A3 runs measured approximately 38.8–39.6 ms/token:
Online decode
A single-device client-observed streaming test used one warm-up followed by five serial 48-token requests:
The online TPOT is measured from first-token arrival to last-token arrival. It includes per-step scheduler, worker IPC, HTTP/SSE, and local client overhead, so it is not directly equivalent to pure kernel or offline TPOT.
Validation
Static and unit checks
git diff --check: passed.Device validation
/healthand/v1/models: passed.Performance test settings:
介绍一下北京故宫temperature=0top_p=1max_new_tokens=48,max_seq_len=256max_new_tokens=512,max_seq_len=1024Related work