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.

Inputs the agent should gather first

Before writing any code, confirm:

1. Layout convention. Two conventions exist and they are not equivalent:
   • GPT-NeoX / HuggingFace-LLaMA: rotate halves. The vector is split into a first half x_lo = x[..., :D/2] and a second half x_hi = x[..., D/2:]. The rotated component is concat(-x_hi, x_lo). Cos/sin tables have shape [max_seq_len, D/2] and are broadcast across the two halves (i.e., cos is duplicated to length D as concat(cos, cos)).
   • GPT-J / original RoPE: rotate adjacent pairs. For each i, the pair (x[2i], x[2i+1]) is rotated together. FlashAttention-2 uses this layout. Picking the wrong layout silently produces broken model output. Confirm by reading the reference model code, not by guessing from the dtype or shape.
2. Head dimension D. Must be even. If the model uses partial RoPE (rotate only the first rot_dim dimensions, leave D - rot_dim untouched), get rot_dim explicitly.
3. Number of heads for Q (H_q) and K (H_kv). Often H_kv < H_q for GQA — the kernel must handle both head counts independently (they share the cos/sin table but loop over their own head range).
4. Cos/sin table dtype. Should be fp32 at table-build time. If the framework hands you fp16 cos/sin tables, flag it — long-position precision loss is a real bug at sequence lengths > 8K.
5. Q/K dtype. Usually fp16 or bf16. The kernel reads in this dtype, casts to fp32 for the rotation, and stores back in the original dtype.
6. Position layout. Two cases:
   • Contiguous prefill: positions [0, 1, 2, ..., N-1] for each sequence, sequences padded to a uniform length.
   • Continuous batching: a positions array of shape [total_tokens] giving each token's true position in its request's sequence. Token i may be position 5 in request A while token i+1 is position 137 in request B.
7. Partial RoPE ratio. Some models (e.g., certain DeepSeek variants) rotate only a fraction of D. If rot_dim < D, the kernel must leave the tail untouched.
8. In-place vs out-of-place. HuggingFace and most serving stacks apply RoPE in place to save memory. Confirm whether the caller expects Q and K to be overwritten or written to separate output buffers.
9. Whether Q and K share a launch. A combined kernel applies RoPE to both in one launch (avoids two launches' overhead but slightly complicates pointer arithmetic). Two separate launches is simpler and only marginally slower for typical shapes.

Required reasoning process

1. Pin down the layout first, then the formula. Write out the rotation in the chosen layout before touching the kernel:

   • NeoX (rotate halves): out_lo = x_lo * cos - x_hi * sin, out_hi = x_hi * cos + x_lo * sin, where cos, sin have shape [D/2] and x_lo, x_hi are each shape [D/2].
   • GPT-J (rotate pairs): for each i, out[2i] = x[2i]*cos[i] - x[2i+1]*sin[i], out[2i+1] = x[2i+1]*cos[i] + x[2i]*cos[i] — wait: out[2i+1] = x[2i+1]*cos[i] + x[2i]*sin[i]. Write this out and double-check signs against a reference implementation before coding.

2. Choose the program decomposition. A typical mapping is one program per (token, head) pair:

   grid = (total_tokens, num_heads)
   pid_t = tl.program_id(0)   # token index in the flat (B*N) layout
   pid_h = tl.program_id(1)   # head index for this Q or K tensor
   

   For continuous batching, total_tokens is the sum of all request lengths, not B * N.

3. Look up the position for this token. Load pos = tl.load(positions_ptr + pid_t). Do not compute pos = pid_t % seq_len — this is wrong for continuous batching and silently produces incorrect positional information for all tokens after the first request.

4. Load the cos/sin row for this position. The cos/sin tables have shape [max_seq_len, D/2]. Compute the row pointer:

   d_offsets = tl.arange(0, BLOCK_D_HALF)         # BLOCK_D_HALF >= D/2
   d_mask    = d_offsets < (rot_dim // 2)         # for partial RoPE
   cos_row = tl.load(cos_ptr + pos * (D // 2) + d_offsets, mask=d_mask, other=1.0)
   sin_row = tl.load(sin_ptr + pos * (D // 2) + d_offsets, mask=d_mask, other=0.0)
   

   other=1.0 for cos and other=0.0 for sin: at masked-out (non-rotated) positions, the rotation reduces to identity, which is what partial RoPE wants.

5. Load the activation halves (NeoX layout).

   x_ptr_base = X + pid_t * stride_token + pid_h * stride_head
   x_lo = tl.load(x_ptr_base + d_offsets,                mask=d_mask, other=0.0).to(tl.float32)
   x_hi = tl.load(x_ptr_base + d_offsets + (D // 2),     mask=d_mask, other=0.0).to(tl.float32)
   

   Cast to fp32 immediately. The rotation arithmetic must be fp32; doing it in fp16 loses precision at large positions (sin/cos of large angles are tiny).

6. Apply the rotation in fp32.

   out_lo = x_lo * cos_row - x_hi * sin_row
   out_hi = x_hi * cos_row + x_lo * sin_row
   

   Sign convention: out_lo = x_lo*cos - x_hi*sin. If you write + x_hi*sin instead of - x_hi*sin, you have rotated in the wrong direction — the model will produce nonsense early-token attention but may not crash.

7. Store results back. Cast to the activation dtype, apply the same boundary mask:

   tl.store(out_ptr_base + d_offsets,            out_lo.to(x_dtype), mask=d_mask)
   tl.store(out_ptr_base + d_offsets + (D // 2), out_hi.to(x_dtype), mask=d_mask)
   

   For partial RoPE, the unrotated tail (d_offsets >= rot_dim/2) must either be skipped (in-place: leave as-is, do not store) or copied through (out-of-place: explicit copy load/store).

8. GPT-J layout variant. Replace step 5 with strided loads:

   x_even = tl.load(x_ptr_base + 2 * d_offsets,     mask=d_mask, other=0.0).to(tl.float32)
   x_odd  = tl.load(x_ptr_base + 2 * d_offsets + 1, mask=d_mask, other=0.0).to(tl.float32)
   

   And step 6 produces out_even, out_odd stored back at the corresponding strided offsets. The cos/sin table indexing is the same.

9. Apply to Q and K, not V. RoPE rotates Q and K. V is left untouched. If the kernel takes a single tensor pointer, the launcher calls it once for Q and once for K. If the kernel takes both Q and K pointers and processes them together, ensure V is never passed.

10. Handle the head-count mismatch for GQA. Q has H_q heads, K has H_kv heads, where often H_kv = H_q / 4 or H_kv = H_q / 8. They share the cos/sin tables (same position, same D). The grid for the Q launch is (total_tokens, H_q); for K it is (total_tokens, H_kv).

Kernel design rules

• BLOCK_D_HALF must be a power of 2, declared tl.constexpr, and at least D / 2. For typical D = 128, use BLOCK_D_HALF = 64. For D = 64, use BLOCK_D_HALF = 32.
• All rotation arithmetic in fp32. Cast inputs to fp32 right after the load; downcast only at the final store.
• The cos/sin table must be loaded as fp32. If the table is stored as fp32 in memory, load directly. If the framework provides fp16/bf16 tables, either accept the precision loss (document it) or cast inside the kernel — but the underlying table-build must be fp32 to avoid loss at large positions.
• positions is an int32 or int64 tensor of shape [total_tokens]. Pass its dtype consistently. Do not compute positions inside the kernel from tl.program_id(0) — that breaks continuous batching.
• All strides (token stride, head stride, D-dimension stride for Q/K; row stride for cos/sin) must be passed as kernel arguments. Do not assume contiguous (total_tokens, H, D) layout.
• For partial RoPE, the rotated portion must be the first rot_dim dimensions in the standard convention. Confirm against the reference model — some research models rotate the last rot_dim instead.
• In-place writes are valid because each (token, head, dim) element is read once and written once with no cross-element dependence beyond the (x_lo, x_hi) pair, which is loaded fully before any store. There is no race within a program.
• A combined Q+K kernel (one launch handles both tensors) saves a launch's worth of latency for decode (length-1 Q/K), which matters in tight serving loops. For prefill, two launches are fine.

Correctness requirements

• Layout must match the reference model's layout exactly. The standard test: feed a known input through the reference (HuggingFace apply_rotary_pos_emb for NeoX-layout models, the original RoPE paper code for GPT-J-layout models) and your kernel; outputs must match within fp16 tolerance (atol=1e-3).
• Sign convention: out_lo = x_lo * cos - x_hi * sin. If swapped to +, the rotation goes the wrong direction. Test against the reference, not against a hand-derived formula.
• The position used for token i must come from the positions array, not from i % seq_len. Continuous batching breaks the latter silently.
• For partial RoPE, the unrotated tail must be unchanged. Verify by checking that out[..., rot_dim:] == x[..., rot_dim:] exactly (bit-identical for in-place; equal for out-of-place copy).
• V must not be passed to this kernel. Apply RoPE to Q and K only.
• The cos/sin table must be built with theta_i = base ** (-2i / D) for i in [0, D/2), then angle = position * theta_i, then cos = cos(angle), sin = sin(angle). Mismatch in the base (10000 vs 1000000 vs an NTK-scaled value) silently changes the embedding and breaks pretrained models.
• The boundary mask (d_mask) must be applied identically on load and store. Asymmetric masking corrupts memory or produces partial outputs.
• Cos/sin computation must be fp32 at table-build time. Storing the table in fp16 loses precision at long sequence positions because sin(p * theta_i) for large p and small theta_i produces values whose fp16 representation is much coarser than the true rotation.

Performance requirements

The agent must reason about:

• Memory traffic. RoPE reads Q and K once and writes them back. For Q of shape (total_tokens, H_q, D) in fp16 and partial RoPE on the full D, the traffic is 2 * total_tokens * H_q * D * 2 bytes (read + write). The kernel is fully memory-bound; no amount of math optimization helps.
• Cos/sin table reuse. Every (token, head) program for the same token loads the same cos/sin row. With H_q + H_kv heads per token, this is read H_q + H_kv times per token. L2 cache absorbs this — do not preload the table into shared memory manually, Triton's tile-level caching handles it.
• Decode-path sensitivity. For decode (length-1 Q/K), the kernel processes B tokens (one per request). Launch latency dominates execution time; aim for one launch handling both Q and K, or fuse RoPE into the attention kernel via FA3-style fusion.
• Combined vs separate launches. A single kernel processing both Q and K in one launch saves one launch's worth of latency (~5-10 μs on H100). For decode, this matters; for prefill, the kernel runtime dwarfs launch overhead and either approach works.
• Head dimension coverage. With D = 128 and BLOCK_D_HALF = 64, the kernel loads D/2 = 64 elements in one tile — fits easily in registers, no inner loop needed. For unusually large head dims (D = 256 on some models), BLOCK_D_HALF = 128 still fits, but verify register pressure does not exceed occupancy targets.
• Do not write a custom RoPE kernel if FA3 fuses it. FlashAttention-3 fuses RoPE into the attention kernel, eliminating a full read-write round-trip on Q and K. A standalone Triton RoPE always loses to this fusion for prefill on H100.

Output format

The agent should produce:

1. The Triton kernel function with @triton.jit, taking pointers (X or Q+K, output, cos, sin, positions), strides for each tensor's token/head/dim axes, shape parameters (D, rot_dim, H_q and/or H_kv), and BLOCK_D_HALF: tl.constexpr.
2. A Python launcher that computes the grid as (total_tokens, H), validates that D is even and that rot_dim <= D, extracts strides, and invokes the kernel separately for Q and K (or jointly, if combined).
3. An explicit comment naming the layout (NeoX or GPT-J) at the top of the kernel. This single comment prevents 90% of layout-confusion bugs in downstream usage.
4. A correctness test comparing against transformers.models.llama.modeling_llama.apply_rotary_pos_emb (for NeoX layout) at fp16 with torch.allclose(atol=1e-3). For GPT-J layout, compare against a reference implementation derived from the original RoPE paper.
5. A continuous-batching test with at least two requests of different lengths and non-zero starting positions, verifying that each token's rotation uses its correct position.
6. Documentation of any partial-RoPE handling, including which dimensions are rotated and which are passed through.

Common failure modes

• Wrong layout (NeoX vs GPT-J). The most insidious bug. The kernel runs, no NaN, no shape mismatch — but the model produces garbled or off-distribution output. Tests against a reference model in the right layout are the only reliable check.
• Wrong sign of sin. out_lo = x_lo * cos + x_hi * sin instead of - x_hi * sin rotates in the opposite direction. The model often still produces plausible-looking text early but degrades catastrophically with longer context. Catches: explicit sign comparison against reference at multiple positions, not just position 0.
• Applying RoPE to V. V must remain unrotated. Applying RoPE to V destroys the value semantics; the model output collapses to gibberish.
• Using token index instead of true position. With continuous batching, pid_t is a flat index into a packed batch, not a position within a sequence. Computing the position as pid_t or pid_t % seq_len produces wrong rotations for every token after the first request. Always read from a positions array supplied by the scheduler.
• fp16 cos/sin precision loss at long positions. If the cos/sin table is stored in fp16, then for position p = 8000 and a small theta_i ≈ 1e-4, the angle p * theta_i ≈ 0.8 is fine — but for p = 100000 (extended-context models with YaRN), precision in fp16 is insufficient. Always build the table in fp32 and store it in fp32, or accept the precision loss explicitly for short-context models only.
• Head-count mismatch silently scaling wrong heads. If the launcher passes Q's H_q as the head count for both Q and K (where K should use H_kv), the K kernel reads beyond its valid head range and writes garbage. Pass the correct head count for each tensor.
• Partial RoPE applied to the wrong half. Some models rotate the first rot_dim dimensions; some rotate the last rot_dim. Check the reference model's slicing convention before coding.
• Storing fp32 results to a fp16 buffer without cast. Forgetting out_lo.to(x_dtype) before tl.store either crashes (Triton type-checks) or silently writes the fp32 bit pattern into the fp16 buffer. Always explicit cast at the store boundary.
• Off-by-one in cos/sin row stride. Cos/sin tables have shape [max_seq_len, D/2] — row stride is D/2, not D. Using pos * D reads the wrong row.
• Forgetting that D/2 may not be a power of 2. For D = 192 (used in some research models), D/2 = 96 is not a power of 2. BLOCK_D_HALF must still be a power of 2 (e.g., 128) with masking on the unused tail.

Review checklist

• Layout convention (NeoX or GPT-J) is named explicitly in a kernel-level comment.
• Cos and sin tables are built in fp32 and have shape [max_seq_len, D/2].
• Position for each token is read from the positions array, not derived from program_id.
• All rotation arithmetic is done in fp32 after immediate cast from input dtype.
• Sign convention out_lo = x_lo*cos - x_hi*sin, out_hi = x_hi*cos + x_lo*sin (NeoX) is verified against a reference.
• RoPE is applied to Q and K only — V is never passed to this kernel.
• For GQA, Q uses H_q heads and K uses H_kv heads in their respective launches.
• Partial RoPE (if applicable) leaves the unrotated tail bit-identical.
• BLOCK_D_HALF is a power of 2 and declared tl.constexpr.
• Boundary masking on load and store is symmetric, with other=1.0 for cos and other=0.0 for sin.
• The output cast to(x_dtype) is present immediately before every tl.store.
• Strides for the token, head, and dim axes are kernel arguments, not hardcoded.
• A correctness test against the reference HuggingFace or original-paper implementation passes at fp16 with atol=1e-3.
• A continuous-batching test with at least two requests of differing lengths and non-zero start positions passes.
• No claim of speedup over FlashAttention-3-fused RoPE without a benchmark.