decode-ops-xpu

Fused decode-step glue for LLM inference on Intel GPUs, loadable through kernels: RMSNorm with a fused residual add, in-place RoPE, SiLU-gate, and greedy argmax. The reference baselines are the eager PyTorch equivalents, matched to 1e-6 or bit-exactly.

This is the non-matmul residue that owns a decode step once attention and the linear layers are fast. Left in eager PyTorch it is several launches and several full passes over the hidden state per layer, and on an integrated GPU those passes are the step. Each op here is one kernel that touches the data once.

Measured live on an Intel Iris Xe: RMSNorm at 0.037 ms against eager's 0.100 (2.7x), argmax over a 128,256 vocabulary at 0.246 against 0.355 (1.4x), matching eager to 1e-6 and bit-exactly for argmax.

Usage

from kernels import get_kernel

d = get_kernel("phanerozoic/decode-ops-xpu", version=1, trust_remote_code=True)

h = d.rms_norm(hidden, norm_weight, 1e-6)
d.fused_add_rms_norm(hidden, residual, norm_weight, 1e-6)
d.rope_inplace(q, k, cos, sin)
gated = d.silu_and_mul(mlp_out)
tokens = d.argmax(logits)

layers.RMSNorm and layers.SiluAndMul are exported for use as Hub layer replacements. version selects the release branch; trust_remote_code is required by kernels for publishers without the trusted-publisher mark.

API

rope_inplace takes query as [tokens, n_q_heads, head_dim] and key as [tokens, n_kv_heads, head_dim] with differing head counts, and rotates the pair (i, i + head_dim // 2) using cos/sin of shape [tokens, head_dim // 2].

Method

Normalization and argmax use one work-group per row with a group reduction, so a row of any length is handled without scratch memory. The normalization kernels hold a row in registers between the sum-of-squares pass and the scaling pass, so it is read from memory once; rows wider than the register budget fall back to re-reading. SiLU-gate moves four floats per access where the width allows it. Argmax carries the index alongside the value so the row is read once.

Measured

Against the eager PyTorch equivalent on Intel Iris Xe (Gen12LP, 96 EUs), torch 2.13.0+xpu, 30 iterations after 5 warmup:

These are memory-bound, so the gains come from traffic rather than arithmetic. The table is reproduced by benchmarks/bench.py, which loads this kernel from the Hub and times it against the eager equivalent on the local device.

Validation

Thirteen checks against PyTorch references on Intel Iris Xe, torch 2.13.0+xpu. RMSNorm matches at 4.77e-07 to 9.54e-07 across shapes (4, 512), (2, 7, 1024), (1, 4096), and (3, 333), the last covering a row that is not a multiple of the work-group size. The residual output of fused_add_rms_norm is bit-exact and its normalized output matches at 4.77e-07. RoPE matches at 2.38e-07 for both query and a grouped-query key with a different head count. SiLU-gate matches at 4.77e-07. Argmax is bit-exact against torch.argmax at vocabularies of 1000, 32000, and 7, and a deliberate two-way tie is verified to resolve to the lower index.

Requirements and limits

• float32 or float16; reductions run in float32 regardless, so the float16 path keeps float32 accumulation accuracy. float16 requires aspect::fp16, which is checked before dispatch.
• Intel GPU with a working Level Zero driver and a torch XPU build.
• The ops ship fake implementations, so they trace under torch.compile without breaking the graph; each is verified to compile into a single frame and to match eager output.

References

Zhang and Sennrich, "Root Mean Square Layer Normalization" (2019); Su et al., "RoFormer" (2021); the NeoX rotary layout.

License

Apache-2.0.