Repository navigation
feat(tokamax): add narrow tokamax 0.0.17 TPU kernel - #1307
tengomucho wants to merge 4 commits into
Conversation
Recipes can vendor a second pinned repository with the `vendor` op, which kernel-port only runs when given `--vendor <name>=<checkout>`. Parse the vendor ops like the source op, clone each pin, and pass it on, so such ports can be checked in CI too. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
|
This PR changes source code for the following kernel(s): If you have the kernel-bot build permissions, then trigger |
Port the Pallas Mosaic TPU kernel bodies of Tokamax 0.0.17 with a kernel-port recipe, as a noarch TPU kernel that other TPU kernels can use through kernel-depends. Tokamax's Op, autotuning and benchmarking layers are left out, so the kernel needs nothing beyond JAX at runtime and adds no Python dependencies. Each op is exposed as its kernel module: ragged dot (including qwix quantized operands), gmm_v2/tgmm_v2, splash and ring attention, and the kernels that were already JAX-only (topk, router_topk, ragged_gather, ragged_gather_reduce, deepseek_rope, csa_gather, sparse_mla, masked_dense_mla, proj_and_save_state, grouped_lowrank_proj, dsv4_compressor, mla, mla_swa, fused_fp8_matmul, fused_moe_decode, mhc). The few parts of qwix (QArray, quantize, dequantize) and immutabledict that the kernels use are vendored; qwix's flax pytree registration is replaced with jax.tree_util.register_dataclass. The kernels_ci tests check that the kernel loads without the dropped dependencies and compare ragged dot, gmm_v2 and splash attention with references on TPU. The kernel is also added to the maintainers registry and to the freshness check's source mapping. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
d0d7d28 to
fab8461
Compare
|
/kernel-bot build tokamax |
|
Build request processed. Command: Dispatched (1):
|
| LOAD_WITHOUT_DROPPED = """ | ||
| import sys | ||
|
|
||
|
|
||
| class Block: | ||
| def find_spec(self, name, path=None, target=None): | ||
| if name.split(".")[0] in {dropped!r}: | ||
| raise ModuleNotFoundError(f"{{name}} is blocked", name=name) | ||
| return None | ||
|
|
||
|
|
||
| sys.meta_path.insert(0, Block()) | ||
|
|
||
| import kernels | ||
|
|
||
| tokamax = kernels.get_kernel("kernels-community/tokamax", version=1) | ||
| print(tokamax.__version__) | ||
| """ |
There was a problem hiding this comment.
this string and executable is a bit strange - can we handle in a different way?
There was a problem hiding this comment.
Yes, agreed and removed. A dropped dependency can only come back with an upstream sync, so the check now lives in the port recipe: port.kdl fails the port if any of them is imported (I added einshape and xprof to the existing checks).
drbh
left a comment
drbh
left a comment
There was a problem hiding this comment.
looks good just left some tiny nits and will approve once resolved.
also I'm looking forward to revisiting this port file soon after the kernel-port tool is updated since I think we can make it a bit more concise (but that will be in a follow up pr)
The dropped-dependency test ran a code string in a fresh interpreter. The port recipe already rejects pydantic, jaxtyping, absl, flax and qwix imports; add einshape and xprof so it covers every dropped dependency, and drop the test. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Summary
Adds
kernels-community/tokamax: the Pallas Mosaic TPU kernel bodies ofTokamax 0.0.17 as a noarch TPU kernel,
so that TPU kernels can use them through
kernel-depends. For instance,ragged-dot-tpu(#1208) currently needs an undeclared tokamax install.This follows the review of huggingface/kernels#972: no new Python
dependencies. Tokamax's
Op, autotuning and benchmarking layers needpydantic, jaxtyping, absl and more at import time, so they are not ported.
The kernel needs nothing beyond JAX at runtime.
tokamax/port/port.kdl), likeeinops.
tokamax/srcis not edited by hand.tokamax.ragged_dot.gmm,tokamax.gmm_v2.gmm_v2):gmm_v2/tgmm_v2;ragged_gather_reduce, deepseek_rope, csa_gather, sparse_mla,
masked_dense_mla, proj_and_save_state, grouped_lowrank_proj,
dsv4_compressor, mla, mla_swa, fused_fp8_matmul, fused_moe_decode, mhc.
QArray,quantizeanddequantize(3 files of qwix 0.1.8).flax.structis replaced withdataclassesplusjax.tree_util.register_dataclass, and the flax nnx registration isdropped.
AsQArray's pytree metadata hashable.count:basemodule (theOplayer). They only used a type alias, and namesre-exported from
reference.disable_multi_core_mode) with thedefault,
False.jax.extendexplicitly. Upstream only gets itthrough modules that are not ported.
expectguards fail the port if an upstream bump brings back pydantic,jaxtyping, absl, flax or a bare qwix import.
scripts/check_ports.py(separate commit) now passes eachvendorop's pinned checkout to kernel-port, so the Check Ports workflow can
reproduce recipes that vendor a second repository.
Testing
On a TPU v6e VM:
scripts/check_ports.py tokamax einops: both regenerate byte for byte.nix build .#bundle(torch-tpuvariant): passes, including theget_kernelimport check.pytest -m kernels_ci tokamax/src/testswithLOCAL_KERNELS, in a venvwithout tokamax, qwix, flax, pydantic or jaxtyping: 7 passed in 21s.
QArrayandAsQArray),gmm_v2and causal splash attention with references.several tilings, eager and
jit).kernel-depends = [{ repo-id = "kernels-community/tokamax", version = 1 }]got the kernel through
get_kernel_dep. It rangmm_v2on Torch TPUtensors through
jax_opand matched a Torch reference.Not tested:
kernels.lockand Nix build:lock-kernel-depsresolvesdependencies on the Hub, so this needs tokamax to be published first.
gmm_v2and splash attention are checked.The other modules are covered by the import and public API tests.
🤖 Generated with Claude Code