fix(kb): pack a rebind overlay's implementation, not just its manifest - #433
Merged
Conversation
_pack_overlay carried the manifest, sitecustomize.py and the _patched/ tree. That is the whole overlay only for a patch-style rebind, where the replacement module IS the _patched/ copy. A `rebinds` overlay names a module that sits at the overlay root, so nothing under _patched/ matched it and the tarball shipped a manifest pointing at a file that was not in the archive. Unpack it, put it on PYTHONPATH, and sitecustomize.py raises ImportError on the first line that matters. Measured on the DeepSeek-V4-Pro authored sparse-MLA decode win: the overlay packed as 2 entries, both metadata. The 30 KB kernel, its seam and the retuned prefill launcher were all absent, so the record promised a reproducible run and shipped nothing to reproduce it with. Naming the manifest's module is still not enough, because that module is regularly a shim. GLM-5.2's dsa_engage_c0_triton is 1.6 KB re-exporting a 23 KB authored kernel from a sibling; gpt-oss-120b's _fwd_kernel keeps its entire HIP source in geak_hip_extend/. So _rebind_modules walks the transitive closure from each rebinds[*].impl_module, and _sibling_imports reads that closure out of the source text with a regex rather than by importing the module -- importing an overlay module on the write path would run its top level, which for an authored kernel means compiling Triton against whatever GPU the writer happens to be sitting on. Names that could reach outside the overlay root -- dotted, absolute, or containing a separator -- are refused rather than resolved.
Seeding the closure from the manifest alone left the other entry point out. An overlay that captures shapes or traces a seam does that work in siblings sitecustomize.py imports directly, with no manifest entry naming them: the sparse-MLA capture overlays import capture_shapes and seam_trace, and both were absent from the tarball while the hook that imports them was in it. Both roots now seed the same walk. The comment claiming a dotted impl_module is refused described something the code does not do -- the name is split on the first dot before that check, so `geak_authored.gemm_flydsl` correctly resolves to the package at the overlay root and is packed whole. Corrected to say so. What is actually refused is a name that is absolute or carries a separator. Five tests cover the closure: the plain rebind, the shim indirection, a submodule rebind, sitecustomize.py's siblings, and the escape refusal. All four of the positive ones fail against the packer this branch replaces. Swept over the 42 overlay directories of the hl_matrix_0824 matrix: 52 files that the previous packer dropped are now carried, no overlay lost a file, and no packed module is left importing a sibling that is not in the archive.
chao-xu-spec
approved these changes
Aug 26, 2026
zihaoanllm
approved these changes
Aug 26, 2026
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
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.
What breaks today
_pack_overlaycarries the manifest,sitecustomize.py, and the_patched/tree. That is thewhole overlay only for a patch-style rebind, where the replacement module is the
_patched/copy.
A
rebindsoverlay leaves the upstream module alone and swaps one symbol forimpl_module's. That module sits at the overlay root, not under_patched/, so nothingmatched it and the tarball shipped a manifest pointing at a file that was not in the archive.
Unpack it, put it on
PYTHONPATH, andsitecustomize.pyraisesImportErroron the first linethat matters.
Measured on the DeepSeek-V4-Pro authored sparse-MLA decode win (2.311x isolated, +4.141% e2e):
Two entries, both metadata. The 30 KB kernel, its seam, and the retuned prefill launcher were all
absent — the record promised a reproducible run and shipped nothing to reproduce it with.
Why naming the manifest's module is not enough
That module is regularly a shim:
dsa_engage_c0_tritonis 1.6 KB offrom dsa_authored_c0_triton import tilelang_sparse_fwdwrapping a 23 KB authored Triton kernel._fwd_kernelwin keeps its entire HIP source ingeak_hip_extend/.So
_overlay_moduleswalks the transitive closure from eachrebinds[*].impl_module, and_sibling_importsreads that closure out of the source text with a regex rather than byimporting the module. Importing an overlay module on the write path would run its top level, which
for an authored kernel means compiling Triton against whatever GPU the writer happens to be
sitting on.
A name that resolves to nothing at the overlay root is an ordinary third-party import and is left
alone, so the closure terminates at the overlay boundary. A submodule rebind
(
geak_authored.gemm_flydsl) is addressed by its top-level package, which is what sits at theroot. Names that are absolute or carry a path separator are refused rather than resolved into
somebody's filesystem.
Second commit: the other entry point
Seeding the closure from the manifest alone left
sitecustomize.pyout. An overlay that capturesshapes or traces a seam does that work in siblings
sitecustomize.pyimports directly, withno manifest entry naming them — the sparse-MLA capture overlays import
capture_shapesandseam_trace, and both were missing from the tarball while the hook that imports them was in it.Both roots now seed the same walk.
The second commit also corrects a comment that claimed a dotted
impl_moduleis refused. The codesplits on the first dot before that check, so it never was; what is actually refused is an
absolute or separator-bearing name.
Verification
Swept
_pack_overlayold-vs-new across the 42 overlay directories of thehl_matrix_0824matrix (31
rebinds, 11 patch-style):old ⊆ new)The last row is an independent closure check — it re-derives each tarball's import graph from the
source and asserts every sibling-resolvable name is present, rather than trusting the packer's own
notion of the closure.
Representative gains:
overlay_c0_triton2 → 5 entries (hl_sparse_decode.py,hl_sparse_decode_seam.py,hl_prefill_ragged_tuned.py);overlay_c1_flydslpicks up the wholegeak_authored/package; the llama-3.1-8B tuning overlays pick upgeak_fp8_rowwise.py.Tests: the three existing
_pack_overlaytests still pass, plus five new ones covering the plainrebind, the shim indirection, a submodule rebind,
sitecustomize.py's siblings, and the escaperefusal. All four positive new tests fail against
origin/main's packer (verified by swappingthe implementation under the test module).