fix(thirdparty/triton_kernels): tolerate upstream module rename#277
Closed
qywu wants to merge 1 commit into
Closed
fix(thirdparty/triton_kernels): tolerate upstream module rename#277qywu wants to merge 1 commit into
qywu wants to merge 1 commit into
Conversation
Older releases of upstream openai/triton's ``triton_kernels`` package shipped ``matmul`` and ``matmul_details``; newer releases (currently the one resolved when ``pip install triton_kernels`` runs against PyPI's ``triton_kernels==0.1.0``, and other channels' newer 1.x builds) rename them to ``matmul_ogs`` / ``matmul_ogs_details``. The unconditional pre-imports in this shim's ``__init__`` raised ``ModuleNotFoundError: No module named 'triton_kernels.matmul'`` whenever the installed package only had the new layout, breaking every downstream import of ``tokenspeed.runtime`` even though the actual consumers in ``tokenspeed_kernel.ops.moe.triton_kernels`` already guard their usages with ``try/except ImportError`` fallbacks. Wrap each submodule import individually so a missing one no longer breaks the chain. The list covers both naming conventions so either flavour of the installed package works. Signed-off-by: Qingyang Wu <willqywu@gmail.com>
This was referenced May 27, 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
1 participant
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
tokenspeed_kernel/thirdparty/triton_kernels/__init__.pypre-imports a list oftriton_kernelssubmodules under the redirected-triton context. The pre-imports were unconditional, so any naming difference between the locally-installedtriton_kernelsand the names hard-coded here broke every downstreamtokenspeed.runtimeimport with:Upstream
openai/triton'striton_kernelsrenamedmatmul→matmul_ogs(andmatmul_details→matmul_ogs_details) at some point between 0.x and 1.x. PyPI'striton_kernels==0.1.0already has the new layout; older builds shipped through other channels still have the old layout.Real call sites in
tokenspeed_kernel/ops/moe/triton_kernels.pyalready guard their actual usages withtry/except ImportError, so the only thing actually broken was the eager pre-import.Change
Wrap each
__import__individually. List both naming conventions so either flavour of the installed package works.Test plan
python -c "from tokenspeed.runtime.entrypoints.engine import Engine"succeeds against the new layout (triton_kernels==0.1.0PyPI build, hasmatmul_ogs). ✅ verified locally.matmul(will still work via the first half of the alternation list).triton_kernels.matmul_ogsis unaffected (no behaviour change to call sites, only to the pre-import shim).Related
Encountered while verifying #272 / #273 / #274 / #275 end-to-end on H100.