refactor: compile L3 host wrappers via signature mode, drop dummy args - #146
Merged
Merged
Conversation
📝 WalkthroughWalkthroughQwen kernel caching now fingerprints PyPTO kernel sources and static module metadata. NPU executor compilation uses a shared JIT path with annotated host wrappers. The serving-owned dispatch module and dummy tensor inputs are removed. Tests use lightweight kernel metadata objects. ChangesQwen kernel compilation and caching
Estimated code review effort: 3 (Moderate) | ~20 minutes Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 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 |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 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 `@tests/unit/model/qwen/test_kernel_cache.py`:
- Around line 93-101: Parameterize
test_params_fingerprint_tracks_every_distinguishing_dimension over
_QWEN3_14B_FINGERPRINT_CONSTS, varying each constant from its baseline value and
asserting _pf produces a different fingerprint. Preserve the existing platform
and kernel-name assertions so every listed static fingerprint constant is
covered.
🪄 Autofix
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 Plus
Run ID: 5fa68a5f-7c95-4849-bec5-a3961b8c8c39
📒 Files selected for processing (7)
pypto-libpypto_serving/model/qwen/kernel_cache.pypypto_serving/model/qwen/npu_executor.pypypto_serving/model/qwen/qwen3_l3_dispatch.pytests/unit/model/qwen/test_kernel_cache.pytests/unit/model/qwen/test_npu_executor.pytests/unit/model/qwen/test_npu_runner_inputs.py
💤 Files with no reviewable changes (2)
- pypto_serving/model/qwen/qwen3_l3_dispatch.py
- tests/unit/model/qwen/test_npu_runner_inputs.py
Comment on lines
+93
to
+101
| def test_params_fingerprint_tracks_every_distinguishing_dimension(): | ||
| args = [_FakeTensor((16, 512), "bfloat16")] | ||
| base = _pf("decode_fwd", args) | ||
| # max_seq (shape) change -> the exact bug the name-only key missed | ||
| assert base != _pf("decode_fwd", [_FakeTensor((16, 2048), "bfloat16")]) | ||
| # dtype, platform, and kernel name all distinguish a binary | ||
| assert base != _pf("decode_fwd", [_FakeTensor((16, 512), "float32")]) | ||
| assert base != _pf("decode_fwd", args, platform="a2a3sim") | ||
| assert base != _pf("prefill_fwd", args) | ||
| base = _pf("decode_fwd", types.SimpleNamespace(MAX_SEQ=2048, VOCAB=512)) | ||
| # MAX_SEQ change -> different binary (the exact bug a name-only key missed) | ||
| assert base != _pf("decode_fwd", types.SimpleNamespace(MAX_SEQ=4096, VOCAB=512)) | ||
| # VOCAB change -> different binary | ||
| assert base != _pf("decode_fwd", types.SimpleNamespace(MAX_SEQ=2048, VOCAB=152064)) | ||
| # platform and kernel name also distinguish a binary | ||
| assert base != _pf("decode_fwd", types.SimpleNamespace(MAX_SEQ=2048, VOCAB=512), platform="a2a3sim") | ||
| assert base != _pf("prefill_fwd", types.SimpleNamespace(MAX_SEQ=2048, VOCAB=512)) |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟡 Minor | ⚡ Quick win
Cover every static fingerprint constant.
Lines 94-101 change only MAX_SEQ and VOCAB. A regression that omits BATCH_PAD, NUM_LAYERS, TOPK, or another listed constant will keep this test green and can reuse an incompatible cached binary. Parameterize the test over _QWEN3_14B_FINGERPRINT_CONSTS.
🤖 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/unit/model/qwen/test_kernel_cache.py` around lines 93 - 101,
Parameterize test_params_fingerprint_tracks_every_distinguishing_dimension over
_QWEN3_14B_FINGERPRINT_CONSTS, varying each constant from its baseline value and
asserting _pf produces a different fingerprint. Preserve the existing platform
and kernel-name assertions so every listed static fingerprint constant is
covered.
ndleslx
force-pushed
the
worktree-qwen-host-wrappers
branch
3 times, most recently
from
August 10, 2026 09:54
141288a to
6585fe2
Compare
Compile both the Qwen3-14B and DeepSeek-V4 L3 host wrappers in annotation-driven signature mode -- no positional dummy tensors -- reading shapes from the wrappers' pl.Tensor[[...], dtype] annotations. - compiler.py: KernelCompiler.compile drops ``dummy_args`` entirely; it always compiles in signature mode and forwards ``**compile_kwargs`` to ``jit_fn.compile`` (e.g. ``name=pl.RUNTIME`` for runtime scalars). - qwen npu_executor.py: read the host wrappers off the already-loaded kernel modules (qwen3_decode_fwd.qwen3_decode_host, etc.); delete the three per-callable dummy_args builders and the qwen3_l3_dispatch injection/file. - deepseek npu_executor.py: pass pl.RUNTIME for each runtime scalar (mtp_num_tokens / num_tokens) instead of building meta-tensor/ctypes sample args; drop _runtime_scalar_compile_args, _PYPTO_TORCH_DTYPES, and the ctypes/inspect imports. - tests: update the KernelCompiler/qwen/deepseek mocks and tests to the signature-mode + RUNTIME-kwarg API; drop the dummy-args-forwarding test. - Bump the pypto-lib submodule gitlink to b6c34f3 (qwen3_14b @pl.jit.host wrappers #914 + rope-dim fix #916). Rebased onto main (which extracted the shared KernelCompiler and removed fingerprinting -- the cache is name-slotted, so the earlier qwen params- fingerprint rework is dropped). Verified: qwen3 accuracy + serving device guards pass locally; 72 unit tests (kernel_compiler + qwen + deepseek); ruff clean.
ndleslx
force-pushed
the
worktree-qwen-host-wrappers
branch
from
August 10, 2026 12:29
6585fe2 to
b93d345
Compare
ndleslx
changed the title
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
Compile both the Qwen3-14B and DeepSeek-V4 L3 host wrappers in annotation-driven signature mode — no positional dummy tensors — reading shapes from the wrappers'
pl.Tensor[[...], dtype]annotations. DeepSeek's runtime scalars (mtp_num_tokens/num_tokens) are passed aspl.RUNTIME(unspecialized).Rebased onto
main(which extracted the sharedKernelCompilerand removed fingerprinting — the cache is now name-slotted).Changes
compiler.py(KernelCompiler.compile):dummy_argsis removed entirely; it always compiles in signature mode and forwards**compile_kwargstojit_fn.compile(e.g.name=pl.RUNTIME).npu_executor.py: read the host wrappers off the already-loaded kernel modules (qwen3_decode_fwd.qwen3_decode_host, etc.); delete the three per-callabledummy_argsbuilders and theqwen3_l3_dispatchinjection/file.npu_executor.py: passpl.RUNTIMEfor each runtime scalar instead of building meta-tensor/ctypes sample args; drop_runtime_scalar_compile_args,_PYPTO_TORCH_DTYPES, and thectypes/inspectimports.KernelCompiler/qwen/deepseek mocks + tests to the signature-mode +RUNTIME-kwarg API; replace the dummy-args-forwarding test with a kwargs-forwarding one.pypto-libsubmodule gitlink tob6c34f3(qwen3_14b@pl.jit.hostwrappers #914 + rope-dim fix #916).Net: 9 files, +51 / −433.
Verification (all on this branch, pypto-lib
b6c34f3)Device guards:
Unit:
72 passed(kernel_compiler + qwen + deepseek);ruffclean.