Add native Qwen3 14B A8W8 serving path#48
Conversation
|
Important Review skippedAuto incremental reviews are disabled on this repository. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
📝 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 |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
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.
vegetabledoww
force-pushed
the
qwen3-a8w8-stage5-serving-slim
branch
3 times, most recently
from
b09f2cf to
2343c8e
Compare
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
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
vegetabledoww
force-pushed
the
qwen3-a8w8-stage5-serving-slim
branch
6 times, most recently
from
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. |
vegetabledoww
force-pushed
the
qwen3-a8w8-stage5-serving-slim
branch
6 times, most recently
from
77e943b to
e6a73fd
Compare
vegetabledoww
force-pushed
the
qwen3-a8w8-stage5-serving-slim
branch
from
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>
vegetabledoww
force-pushed
the
qwen3-a8w8-stage5-serving-slim
branch
3 times, most recently
from
412c199 to
39924ca
Compare
|
vegetabledoww
force-pushed
the
qwen3-a8w8-stage5-serving-slim
branch
from
39924ca to
2a24e51
Compare
|
Maintainer note: Before merging this PR, please update the |
vegetabledoww
force-pushed
the
qwen3-a8w8-stage5-serving-slim
branch
4 times, most recently
from
1f8d5f5 to
8a2b15f
Compare
Add the compressed-tensors A8W8 loader for Qwen3-14B checkpoints and extend the common PyPTO runtime plumbing for PTO ISA pinning and L2 child-memory tensors.
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.
vegetabledoww
force-pushed
the
qwen3-a8w8-stage5-serving-slim
branch
from
8a2b15f to
91aae68
Compare
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.
1 participant
Summary
Register
qwen3-a8w8as a native Qwen3-14B serving model format and route generation through the normalLLMEngineflow, while keeping the existing BF16 Qwen3-14B path isolated.This PR adds the serving-side pieces needed to run the Qwen3-14B A8W8 kernels from
pypto-lib:DistributedProgramHOST wrappers onlyqwen3_l3_dispatch.pymodule for both model formats, with a closure-based A8W8 wrapper factory that keeps A8W8 kernel dependencies isolated from the BF16 dispatch globalsblock_table,slot_mapping, and INT8 KV scale caches in the serving runner--model-format qwen3-a8w8innpu_generate.pypage_size,max_batch_size,kv_dtype,weight_dtype, andtotal_kv_pagesThe previous A8W8 L2 dispatch fallback has been removed from the active serving path. The current implementation intentionally supports the L3 path only. The existing BF16 HOST wrappers and BF16 executor call path remain unchanged.
The latest A8W8 decode scheduling and attention performance work used by the validation below is tracked in pypto-lib PR #760.
Commit Structure
The PR is split into four commits:
feat(qwen3-a8w8): add loader and L2 runtime supportfeat(qwen3-a8w8): add A8W8 L2 model runnerfeat(qwen3-a8w8): compile kernels and wire generation CLInpu_generate.pyto select BF16 or A8W8 model formats.Use L3 dispatch for Qwen3 A8W8DistributedProgramcallables.qwen3_l3_dispatch.pymodule instead of adding a parallel dispatch module.prefill_fwd/decode_fwdglobals.Validation
Static checks:
python -m py_compile examples/model/qwen3_14b/npu_generate.py examples/model/qwen3_14b/runner/npu_executor.py examples/model/qwen3_14b/runner/npu_executor_a8w8.py examples/model/qwen3_14b/runner/npu_runner_a8w8.py examples/model/qwen3_14b/runner/qwen3_l3_dispatch.pypre-commit run --files examples/model/qwen3_14b/runner/qwen3_l3_dispatch.py examples/model/qwen3_14b/runner/npu_executor_a8w8.pygit diff --checkpytest -q tests/test_batching.py(22 passed, 4 warnings)Hardware generation checks were run with the A8W8 kernels from pypto-lib PR #760.
L3-only quick validation after dispatch consolidation
qwen3-a8w8介绍一下北京故宫7.579s55.1 ms/token18.15 tok/s0L3-only 48-token validation
qwen3-a8w8介绍一下北京故宫7.224s55.3 ms/token18.09 tok/s0Related PRs / Issues