thox-qk-rope
Query-key RMSNorm and NeoX rotary embedding for local decode. Source and a tested numerical contract. Not a Kernel Hub binary.
Your AI. Your Data. Your Rules.
Source git: ttracx/thox-qk-rope at de3baa1. This Hub repository is the release card. The kernel source, tests, and paper/main.pdf are in that git tree.
| Item | State |
|---|---|
| What this repository is | Release card for the kernel source |
| Kernel Hub binary | Not uploaded. get_kernel("Thox-ai/thox-qk-rope") will not resolve until a kernel-builder upload |
| Reference tests | NumPy contract and the fp32 C spec. No device timing |
Backends in build.toml |
CUDA and CPU, float32. Metal is a shader note only |
| Companion kernel | Thox-ai/thox-edge-ops (residual RMSNorm and SiLU-and-mul). This repository does not modify it |
thox-edge-ops covers the pre-norm residual and the SwiGLU gate. It does not apply QK-norm or RoPE. Qwen3 adds per-head QK-norm before RoPE. Qwen2.5 needs RoPE and does not use QK-norm. Those are the two calls this kernel exists to serve. Gemma's (1 + weight) norm is out of scope, as it already is for thox-edge-ops.
Contract
y = weight * dtype(x)(x * rsqrt(mean(x^2) + eps)) # Qwen3RMSNorm promotion
out = rotate_half(y) with the caller-supplied cos, sin
The caller owns the frequency table (theta, ABF, YaRN, partial rotary). This kernel does not invent positions.
Layouts: [..., D] with a broadcastable cos/sin table, or Qwen3's pre-transpose view [batch, seq, heads, dim] with cos/sin [batch, seq, rotary]. rotary is even and may be shorter than dim; the tail is copied after the norm. v1 native kernels are float32.
Tests
python3 -m unittest tests.test_reference tests.test_spec -v
Paper
paper/main.tex in the git repository is the technical note. It reports no speedup. Authors: Craig Ross and Tommy Xaypanya. No arXiv id yet, so there is no Hugging Face paper page to claim.
Apache-2.0. Copyright 2026 THOX.ai LLC.