thox-edge-ops

Fused inference kernels for small-model decode on THOX.ai edge hardware. Measured for v2 on x86_64 hosts and an NVIDIA desktop GPU (RTX 4060 Ti); Jetson Orin, Apple silicon and the Arm device class (ThoxAir, ThoxMini, ThoxClip) are targets: the Metal backend is in the source tree, but v2 has not been run on a Mac and no Metal variant is published for v2, and no aarch64 build is published for v2 (see Validation and Limitations). v2 adds the Gemma weight convention, two gated activations, rotary embedding and a fused per-head QK-norm plus RoPE prologue for Qwen3 and Gemma 3 class models, and a CPU backend that selects AVX2, FMA and F16C at run time from one baseline artifact.

Your AI. Your Data. Your Rules.(TM)

Quickstart

pip install -U kernels torch

import torch
from kernels import get_kernel

ops = get_kernel("Thox-ai/thox-edge-ops", version=2, trust_remote_code=["Thox-ai/thox-edge-ops"])
dev = "cuda" if torch.cuda.is_available() else "cpu"
dt = torch.float16 if dev == "cuda" else torch.bfloat16

x = torch.randn(1, 1152, device=dev, dtype=dt)
w = torch.zeros(1152, device=dev, dtype=dt)                     # Gemma stores zero-centred weights
y = ops.rms_norm(x, w, 1e-6, weight_offset=1.0)                 # (1 + w) convention

residual = torch.randn_like(x)
x, residual = ops.fused_add_rms_norm(x, residual, w, 1e-6, weight_offset=1.0)  # in place

h = ops.gelu_tanh_and_mul(torch.randn(1, 2 * 6912, device=dev, dtype=dt))

# Qwen3 / Gemma 3 attention prologue: per-head q_norm, k_norm, then RoPE, one launch per token
head, max_pos = 256, 4096
inv = 1.0 / (10000.0 ** (torch.arange(0, head, 2, device=dev).float() / head))
freqs = torch.outer(torch.arange(max_pos, device=dev).float(), inv)
cos_sin_cache = torch.cat([freqs.cos(), freqs.sin()], dim=-1).to(dt)
positions = torch.tensor([7], device=dev)
q = torch.randn(1, 8, head, device=dev, dtype=dt)
k = torch.randn(1, 4, head, device=dev, dtype=dt)
q_norm_weight = torch.zeros(head, device=dev, dtype=dt)
k_norm_weight = torch.zeros(head, device=dev, dtype=dt)
ops.fused_qk_norm_rope(positions, q, k, q_norm_weight, k_norm_weight, 1e-6, 1.0, head, cos_sin_cache)

trust_remote_code is required because Thox-ai is not on the kernels default trusted-publisher list; scope it to this repository as shown.

Operations

All functions accept float32, float16 and bfloat16 tensors. Non-contiguous inputs are made contiguous by the Python wrapper. Invalid shapes, mismatched weight sizes and negative epsilon raise ValueError. cos_sin_cache is [max_position, rot_dim] with cos in the first half of each row and sin in the second; a token whose position lies outside the cache is left unrotated (pass check_positions=True to raise instead, at the cost of a device sync).

Layers

All layers are stateless, declare can_torch_compile = True and has_backward = False.

Use with transformers

from kernels import LayerRepository, Mode, kernelize, register_kernel_mapping

repo = LayerRepository(repo_id="Thox-ai/thox-edge-ops", layer_name="GemmaRMSNorm", version=2,
                       trust_remote_code=["Thox-ai/thox-edge-ops"])
register_kernel_mapping({"GemmaRMSNorm": {"cuda": repo, "mps": repo, "cpu": repo}})
model = kernelize(model, mode=Mode.INFERENCE)

Pass use_fallback=False to kernelize to fail loudly instead of keeping the original forward. Mixed-dtype checkpoints (fp32 norm weight with fp16 or bf16 activations) keep transformers' promotion semantics through the reference path.

Backends

Performance

Operator-level latency against eager PyTorch written the way transformers writes it, including Python dispatch. Not model throughput. Full tables, raw samples and the paper are in the source repository.

NVIDIA RTX 4060 Ti (CUDA, WSL2), float16. Geometric mean of median-latency ratios over the shapes below; decode = one token per call. 30 trials x 50 calls per configuration, eager and kernel interleaved.

x86_64 CPU, AVX2 dispatch, float32. Geometric mean of median-latency ratios over the shapes below; decode = one token per call. 30 trials x 50 calls per configuration, eager and kernel interleaved.

At decode sizes the GPU operators are launch bound, so the gain is the number of launches removed (about sixteen to one for the fused prologue). On the CPU host the bfloat16 gated activations at prefill sizes are slower than eager because ATen dispatches AVX-512 with sixteen lanes against the kernel's AVX2 eight; keep those on the decode path or route prefill to eager.

Numerics

• Accumulation is fp32 for every dtype; outputs round once.
• fused_add_rms_norm rounds x + residual to the input dtype before squaring, so the stored residual is bit-identical to eager x + residual (asserted with zero tolerance).
• fused_qk_norm_rope rounds the per-head norm to the input dtype before rotating, matching the Qwen3 and Gemma 3 module order; it is asserted equal to rms_norm followed by rotary_embedding.
• Test tolerances against the fp32 eager reference: 1e-5 (fp32), 2e-3 (fp16), 1.6e-2 (bf16).

torch.compile

Every op registers a fake implementation in the build-unique op namespace, so torch.compile(fullgraph=True) traces them with zero graph breaks.

Validation

Published variants

Branch v2, 23 builds from kernel-builder: CPU for torch 2.10 to 2.15, and CUDA 12.6, 12.8, 13.0, 13.2 and 13.4 builds matching each torch release's supported toolkits, all x86_64-linux. The kernels loader picks the one matching the installed torch and CUDA. Branch v1 keeps the v1 builds, including Metal.

Limitations

• Inference only: no backward pass.
• bfloat16 on Metal runs an eager fallback (Metal 3.0 has no bfloat type).
• fused_add_rms_norm needs distinct, contiguous x and residual; anything else takes an out-of-place path.
• Per-head norm weights are shared across heads ([head_size]), which is the Qwen3 and Gemma 3 layout; per-head-distinct weights are not supported.
• No AVX-512 unit yet; on AVX-512 hosts prefill-size gated activations in bf16 favour eager.
• On CPUs without AVX2 the portable path wins only for the norms and rotary at decode; its gated activations evaluate exp and tanh per element and are slower than eager, so keep SiluAndMul and GeluTanhAndMul on eager there.
• Published Hub variants are x86_64-linux (CPU and CUDA). Metal and aarch64 are source-only in v2.

Offline and pinned deployment

# pyproject.toml
[tool.kernels.dependencies]
"Thox-ai/thox-edge-ops" = 2

kernels lock .        # writes kernels.lock with the resolved revision
kernels download .    # pre-fetches it into the HF cache

Architecture

flowchart LR
  A["Python API: rms_norm, fused_add_rms_norm, silu_and_mul, gelu_tanh_and_mul, rotary_embedding, fused_qk_norm_rope, layers"] --> B["torch.ops namespace, unique per build"]
  B --> C["CUDA kernels (sm_87 and up)"]
  B --> D["Metal kernels (Apple silicon)"]
  B --> E["CPU dispatcher"]
  E --> F["AVX2 + FMA + F16C unit"]
  E --> G["baseline at::vec unit (NEON on aarch64)"]
  A -. "unsupported dtype or device" .-> H["eager fp32 reference"]

Breaking changes ship as a new major version branch. v1 (three ops, no weight offset) remains on branch v1 for anyone who pinned it.

Research

Tommy Xaypanya. "thox-edge-ops v2: fused decode kernels for small models on THOX edge hardware." THOX.ai Research, v2.0.0 (revised public edition), November 5, 2026. Preprint, not peer reviewed.

• Paper page: https://thox.ai/research/papers/thox-edge-ops
• Zenodo record (DOI): https://doi.org/10.5281/zenodo.23971227 (concept DOI: 10.5281/zenodo.23971226), CC BY 4.0
• Whitepapers dataset: https://huggingface.co/datasets/Thox-ai/thox-whitepapers (folder edge-ops/)

The paper's Section 4 places these kernels in the THOX stack (edge hardware, ThoxOS and ThoxInux, MeshStack) and states the hardware qualification matrix this card follows: the x86_64 workstation and hub class is measured for v2; Jetson Orin, Apple silicon and the Arm device class are targets only for v2.

License

Apache-2.0. Copyright 2026 THOX.ai LLC.