rope-hf-neuron-kernels

Surfaces the nki-lib rope_hf kernel (rotary position embedding, HuggingFace [batch, heads, seq, head_dim] layout) on the Hugging Face Kernel Hub, loadable through the standard kernels / get_kernel interface.

This repo does not contain a copy of the kernel. It imports rope_hf from the installed nki-lib package (import name nkilib) and re-exports thin wrap_nki wrappers. The kernel source is maintained in nki-lib; this repo is a discovery/loading shim.

When this kernel helps (read this first)

Benchmarked embedded in a real Qwen2.5-1.5B attention block, whole model compiled with torch.compile(backend="neuron"), against neuronx-cc-fused native PyTorch RoPE (the fair baseline). The verdict depends on sequence length:

Use this kernel for short or long sequence workloads. In the mid-range, the neuronx-cc-fused native RoPE is as fast or faster — the kernel's fixed HOP-boundary cost isn't amortized there. At long S, native rotate_half RoPE's intermediate materialization blows up (e.g. at S=4096 native jumps to 21 ms vs the kernel's 13 ms) and the fused BF16 kernel pulls clearly ahead.

Requirements

• nki — declared in metadata.json python-depends (kernels neuron-backend allow-list; NVIDIA's CUDA equivalent is nvidia-cutlass-dsl).
• nki_library (nkilib) — the actual kernel source, imported at load time. Not on the kernels allow-list, so not declared in metadata; an external runtime requirement. Pre-installed in the AWS PyTorch Native Beta container and on the AWS Neuron pip index (not public PyPI).
• A torch.neuron-registered PyTorch build (PyTorch Native / TorchNeuron Beta 5+; torch.neuron is registered natively there — no shim needed).
• Trainium (trn2 recommended).

Usage

from kernels import get_kernel
rope = get_kernel("jburtoft/rope-hf-neuron-kernels", version=1, trust_remote_code=True)

# functional (paired Q/K), HuggingFace layout [B, H, S, D]:
q_rot, k_rot = rope.rope_hf_pair(q, k, cos, sin, lnc=2)

# or as an nn.Module inside a torch.compile(backend="neuron") model:
rope_mod = rope.RopeHF(lnc=2)
q_rot, k_rot = rope_mod(q, k, cos, sin)

• cos/sin may be [S, D] (shared) or [B, S, D] (per-batch).
• q_heads != k_heads (GQA) is supported.
• Constraint: S must be divisible by 128 * lnc (256 at lnc=2).

Validated results (trn2.3xlarge, LNC=2, PyTorch Native Beta 5)

Qwen2.5-1.5B attention block (RMSNorm + QKV proj + rope + SDPA + o_proj), 4 layers stacked, whole model compiled, BF16, prefill. nki-lib rope_hf (wrap_nki) vs compiled native RoPE, identical weights:

Parity is high throughout (BF16 accumulation drift grows slightly at very long S, still well within tolerance). See examples/.

License

Apache-2.0 (same as nki-lib). The kernel is © Amazon.com, Inc. and distributed as part of nki-lib.