Add Qwen3 14B A8W8 kernels#642
Merged
zhangqi-chen merged 1 commit intoJul 9, 2026
Merged
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:
📝 WalkthroughWalkthroughExtends ChangesGolden Runner Compatibility and save_actual_data
Qwen3-14B A8W8 Decode Kernel
Qwen3-14B A8W8 Prefill Kernel and rms_lm_head Fix
Estimated code review effort🎯 5 (Critical) | ⏱️ ~120 minutes Possibly related PRs
Suggested labels
🚥 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}
Contributor
There was a problem hiding this comment.
Code Review
This pull request introduces full-layer decode and prefill forward passes for Qwen3-14B with A8W8 quantization, alongside compatibility shims and patching utilities in the test runner. The review feedback focuses on performance optimizations in the prefill implementation, suggesting the vectorization of row-by-row dequantization loops (for Q, K, and V projections) and the SiLU activation loop to eliminate scalar loop overhead on the NPU. Additionally, the reviewer recommends removing redundant pre-zeroing of the SiLU tile and simplifying a conditional boolean assignment in the runner utility.
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[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 8
🤖 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 `@golden/runner.py`:
- Line 317: Ruff B010 is flagging the constant-name setattr usage in the Worker
setup. Replace the setattr-based property assignment in the code that defines
chip_contexts, and the similar assignment mentioned for the other location, with
direct class attribute assignment while preserving the same property behavior.
Use the existing Worker and property(_chip_contexts) symbols to locate the
affected spots and update both occurrences consistently.
- Around line 804-805: The docstring for the data-saving behavior is incomplete:
`save_actual_data` is documented only for the `golden_data` path, but `Runner`
also writes `data/actual` in the non-golden path even when it is false. Update
the documentation near the `Runner`/data persistence logic to describe both
branches clearly, or adjust the condition so `data/actual` is only saved when
intended; make sure the explanation matches the behavior in the relevant
`save_actual_data` handling and `golden_data` flow.
- Around line 391-395: The import handling in runner.py is too broad because the
ModuleNotFoundError around rebuild_kernel_cpp_from_pto also hides missing
transitive dependencies inside pypto.runtime.debug.pto_rebuild. Update the
fallback in the code path that imports rebuild_kernel_cpp_from_pto so it only
returns work_dir when the debug module itself is unavailable, and let other
import failures surface instead of silently reusing stale artifacts.
- Around line 1100-1101: run_jit() is missing the same generated-artifact
patching that run() applies, so JIT-compiled or reloaded runtime_dir artifacts
can still fail on bitcast/host_orch compatibility. After
_compile_jit_with_compat() and after any runtime_dir reload logic in run_jit(),
invoke the same patch helpers used by run() on the compiled output_dir /
work_dir before proceeding. Use the existing run(), _compile_jit_with_compat(),
and runtime_dir handling paths as the reference points so the JIT flow matches
the non-JIT artifact normalization.
- Around line 965-975: The JIT fallback compile path in runner.py is dropping
distributed settings, so entries that rely on L3/distributed mode are compiled
with the wrong configuration. Update the ir.compile call in the fallback that
rebuilds the ir.Program to forward compile_cfg["distributed_config"] (or the
equivalent distributed config value) alongside the existing run_config fields,
using the ir.compile entry point and the surrounding fallback logic as the place
to fix it.
In `@models/qwen3/14b/decode_layer_a8w8.py`:
- Line 193: The comments in the affected decode layer sections use ambiguous
multiplication glyphs that Ruff flags as RUF003. Update the marked comment text
in decode_layer_a8w8.py to replace each non-ASCII multiplication symbol with
plain ASCII wording like x or by, and apply the same cleanup in the other
flagged comment blocks near the related SEQ_TILE/head references.
- Around line 19-25: The copied harness docstring in decode_layer_a8w8.py is
stale and conflicts with the current serving path. Update or remove the header
block near the decode harness so it no longer says the program does not compile
and no longer references obsolete standalone script names; instead, describe the
native serving flow through decode_fwd and keep the docstring aligned with the
current kernel entrypoint and behavior.
In `@models/qwen3/14b/prefill_fwd_a8w8.py`:
- Around line 1120-1125: The final writeback in the prefill path is assembling a
full token tile into dynamic output even when the last tile is partial. Update
the writeback logic around out_chunk, out_out_quant_chunk_bf16, and pl.assemble
in prefill_fwd_a8w8 to trim the chunk with the same set_validshape handling used
for resid1_chunk before any out writeback, and make sure the same fix is applied
in the debug branches as well.
🪄 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: 25cbdbf0-d212-4a1a-b664-b3666d82adc1
📒 Files selected for processing (5)
golden/runner.pymodels/qwen3/14b/decode_layer_a8w8.pymodels/qwen3/14b/prefill_fwd_a8w8.pymodels/qwen3/14b/rms_lm_head.pytests/golden/test_runner.py
vegetabledoww
force-pushed
the
qwen3-a8w8-stage4-generate
branch
from
c0e0128 to
c7ada3a
Compare
vegetabledoww
force-pushed
the
qwen3-a8w8-stage4-generate
branch
6 times, most recently
from
e8f46c5 to
f9f1cc6
Compare
vegetabledoww
force-pushed
the
qwen3-a8w8-stage4-generate
branch
9 times, most recently
from
a9d33c9 to
352b23b
Compare
Add Qwen3-14B A8W8 prefill/decode kernels while keeping the path isolated from the BF16 Qwen3-14B implementation. Decode composes generic PyPTO primitives directly for INT8 accumulation dequantization instead of relying on a PyPTO model-specific composite op. Validation: 48-token generation for prompt '介绍一下北京故宫' completed with normal text quality, TTFT 6.220s, TPOT 70.5 ms/token after rebuilding pypto_core.
vegetabledoww
force-pushed
the
qwen3-a8w8-stage4-generate
branch
from
352b23b to
44c6f3c
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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:
prefill_hiddensupport for Qwen3-14B A8W8 hidden-state prefill chunksdecode_fwdlayer kernel used by serving decodepypto-servingrather than carrying standalone lib-side runnersImplementation Notes
The lib-side deliverable is intentionally limited to kernels and kernel-adjacent compatibility:
models/qwen3/14b/prefill_fwd_a8w8.pyimplements the A8W8 prefill hidden pathmodels/qwen3/14b/decode_layer_a8w8.pyimplements the A8W8 decode layer path used by servinggolden/runner.pycompatibility changes keep generated/runtime artifacts runnable for the JIT pathValidation
End-to-end validation was run through the native serving stack with this kernel PR plus the matching serving/backend PRs:
介绍一下北京故宫7.054s67.8 ms/token14.76 tok/sFocused serving and PyPTO checks also passed in the matching PRs.
Related PRs / Issues