write-triton-kv-cache-append-kernel

Skill: Write a Triton KV Cache Append Kernel

Purpose

Guide the agent through implementing a Triton kernel that writes newly computed K and V tensors into a pre-allocated KV cache during LLM inference. This covers two cache layouts (contiguous and paged / vLLM-style PagedAttention), unified prefill and decode handling via a slot_mapping tensor, GQA/MQA where the cache stores fewer heads than Q, optional fp8/int8 quantized KV cache with scaling, coalesced versus scattered write patterns, and the boundary checks required to avoid corrupting other requests' cache regions. The kernel runs once per layer per forward step; correctness is the dominant concern, throughput is secondary because the data volume per call is small.

Use this when

• You are building a custom serving stack and need a KV cache append step that is not provided by vLLM or FlashInfer.
• You need a paged KV cache layout (block_table indirection) with a non-standard block_size, an exotic block layout, or a sharding scheme like RingAttention that the existing libraries do not support.
• You need fp8 or int8 quantized KV cache with per-block or per-tensor scales, and the in-tree kernels do not match your scaling convention.
• You want a unified kernel that handles prefill (many new tokens per request) and decode (one new token per request) through the same slot_mapping interface.
• You are fusing the append into a custom QKV epilogue and need a reference standalone kernel to verify against.

Do not use this when

• You are inside vLLM and the existing reshape_and_cache_flash (or reshape_and_cache) kernel already covers your layout. It is heavily tuned and maintained.
• You are using FlashInfer — its append_paged_kv_cache is well optimized and integrates with the rest of the FlashInfer attention API.
• You only need a contiguous cache and a simple PyTorch indexing assignment is fast enough. For small batch sizes the launch overhead of a custom kernel can exceed cache[:, :, pos, :] = new_k.
• You are tempted to fuse the append into the QKV projection in your first iteration. The projection's tile shape constraints and the cache's scatter pattern rarely align cleanly. Write the standalone kernel first, profile, then consider fusion if the launch overhead actually matters.
• You need a backward pass. KV cache append is inference-only; there is no autograd-relevant version of this kernel.