write-triton-rope-kernel

Skill: Write a Triton RoPE Kernel

Purpose

Guide the agent through implementing a correct Triton kernel that applies Rotary Position Embeddings (RoPE) to query and key tensors before attention. This covers the two incompatible layout conventions (GPT-NeoX/HuggingFace-LLaMA vs GPT-J/original-paper), pre-computed cos/sin table consumption, per-token position handling for continuous batching, partial-RoPE masking, and the precision discipline required to keep cos/sin in fp32 while applying to fp16/bf16 activations. RoPE is the dominant positional encoding in LLaMA, Mistral, Qwen, Gemma, GPT-NeoX, and most decoder-only LLMs trained after 2022, so getting this kernel right is load-bearing for inference correctness.

Use this when

• You are building an inference serving stack (vLLM-style, TGI-style, custom) that does not use FlashAttention-3's fused RoPE-in-attention path, and you need a standalone RoPE op for prefill or decode.
• You need a decode-time RoPE kernel: Q/K of length 1 per request, where the launch overhead of a fused FA3-style attention kernel exceeds the cost of a tiny dedicated RoPE kernel.
• You need a custom RoPE variant (NTK-aware scaling, YaRN, longrope, partial RoPE on the first N dims only) where the framework's stock kernel does not match the model definition.
• You need to support continuous batching where each request has a distinct position offset and the standard contiguous-position kernel cannot be used.
• You are porting a model whose RoPE layout (NeoX vs GPT-J) does not match what your inference framework provides.

Do not use this when

• You are using FlashAttention-3 or a similar fused attention kernel that already applies RoPE inside attention. Adding a separate RoPE pass duplicates work and rotates Q/K twice.
• You are running training or inference where HuggingFace's apply_rotary_pos_emb is fast enough — for non-tight loops the Python-level reference is fine and avoids a custom kernel surface.
• The model uses a different positional encoding (ALiBi, T5 relative bias, learned absolute embeddings). RoPE is not a drop-in substitute.
• The tensor layout is exotic and you have not yet decided whether RoPE is applied to the (B, N, H, D) or (B, H, N, D) view. Resolve layout first; the kernel structure depends on it.