Now let me write the CUDA kernel. The operation is: 1. For each (b, t) pair in the new data: - Get position p = positions[b, t] - Apply RoPE rotation to k_new[b, t] using cos[p] and sin[p] - Store rotated key at updated_k[b, p] - Store v_new[b, t] at updated_v[b, p] 2. Stack and return [updated_k, updated_v] Key optimizations: - Parallelize over (batch, update_len, num_heads) - that's 2*16*4 = 128 elements for this problem, but could be larger - Coalesced memory access - Use shared memory for cos/sin if beneficial (but they're small - 128*16 = 2048 floats = 8KB) - Vectorized loads/stores where possible Let me write the kernel: ```cpp #include <torch/extension.h> #include <cuda.h> #include <cuda_runtime.h> template <typename scalar_t> __global__ void rope_update_kernel( scalar_t* __restrict__ updated_k, scalar_t* __restrict__ updated_v, const scalar_t* __restrict__ k_new, const scalar_t* __restrict__ v_new, const scalar_t* __restrict__ cos, const scalar_t* __restrict__ sin, const int* __restrict__ positions, const scalar_t* __restrict__ cache_k, const scalar_t* __restrict__ cache_v, int batch_size, int update_len, int num_heads, int head_dim, int cache_len ) { // Grid: (batch_size * update_len * num_heads) blocks // Each block processes one (b, t, h) combination int idx = blockIdx.x; int h = idx % num_heads; idx /= num_heads; int t = idx % update_len; int b = idx / update_len; if (b >= batch_size) return; int pos = positions[b * update_len + t]; // Copy from cache first for (int d = 0; d < head_dim; d++) { int cache_idx = b * cache_len * num_heads * head_dim + pos * num_heads * head_dim + h * head_dim + d; int new_idx = b * cache_len * num_heads * head_dim + pos * num_heads * head_dim + h * head_dim + d; updated_k[new_idx] = cache_k[cache_idx]; updated_v[new_idx] = cache_v[cache_idx]; } // Apply RoPE rotation to key // Rotated: even = x_even * c - x_odd * s, odd = x_even * s + x_odd * c int new_key_offset = b * cache_len * num_heads * head_dim + pos * num_heads * head_dim + h * head_dim; int input_key_offset = b * update_len * num_heads * head_dim + t * num_heads * head_dim + h * head_dim; int cos_offset = pos * (head_dim / 2); for (int d = 0; d < head_dim / 2; d++) { scalar_t c = cos[cos_offset + d]; scalar_t s = sin[cos_offset + d]; scalar_t x_even = k_new[input_key_offset + d * 2]; scalar_t x_odd = k_new[input_key_offset + d * 2 + 1]; updated_k[new_key_offset + d * 2] = x_even * c - x_odd * s; updated_k[new_key_offset + d * 2 + 1] = x_even * s + x_odd * c; } // Copy new value int new_val_offset = b * cache_len * num_heads * head_dim + pos * num_heads * head_dim + h * head_dim; int input_val_offset = b * update_len * num_heads * head_dim + t * num_heads * head_dim + h * head_dim; for (int d = 0; d < head_dim; d++) { updated_v[new_val_offset + d] = v_new[input_val_offset + d]; } } ``` Wait, I need to think about this more carefully. The original code: 1. Creates updated_k and updated_v as clones of the cache 2. For each (b, t), gets position p, rotates k_new[b,t] and stores at updated_k[b, p], also stores v_new[b,t] at updated_v[b, p] 3. Returns stack([updated_k, updated_v], dim=0) So the output shape is [2, batch, cache_len, num_heads, head_dim]. Let me reconsider the kernel design. I should: 1. First copy cache_k to updated_k and cache_v to updated_v (or handle this in the kernel) 2. Then for each (b, t, h), apply RoPE and update Actually, let me think about the parallelization more carefully. The cache copy is a large operation (batch * cache_len * num_heads * head_dim elements). But we only need to copy cache_k and cache_v to their updated counterparts. A simple approach: - Kernel 1: Copy cache_k to updated_k, cache_v to updated_v - Kernel 2: Apply RoPE rotation and store the new values Or we can do it in one kernel where each thread handles a specific (b, pos) element and applies updates if that position matches any positions[b, t]. Actually, the simplest approach for correctness is: 1. Launch a kernel to copy cache_k -> updated_k, cache_v -> updated_v 2. Launch another kernel to apply RoPE and store new values Let me write this more carefully: ```cpp // Kernel to copy cache and apply updates template <typename scalar_t> __global__ void rope_update_kernel( scalar_t* __restrict__ updated_k, scalar_t* __restrict__ updated_v, const scalar_t* __restrict__ k_new, const scalar_t* __restrict__ v_new, const scalar_t* __restrict__ cos, const scalar_t* __restrict__ sin, const int* __restrict__ positions, const scalar_t* __restrict__ cache_k, const scalar_t* __restrict__ cache_v, int batch_size, int update_len, int num_heads, int head_dim, int cache_len, int total_elements // batch_size * cache_len * num_heads * head_dim ) { int idx = blockIdx.x * blockDim.x + threadIdx.x; if (idx >= total_elements) return; // Calculate indices int d = idx % head_dim; idx /= head_dim; int h = idx % num_heads; idx /= num_heads; int pos = idx % cache_len; idx /= cache_len; int b = idx; // Copy from cache int cache_offset = b * cache_len * num_heads * head_dim + pos * num_heads * head_dim + h * head_dim + d; int out_offset = cache_offset; // Same layout updated_k[out_offset] = cache_k[cache_offset]; updated_v[out_offset] = cache_v[cache_offset]; } ``` Wait, this is getting complicated because we need to also apply the RoPE updates. The issue is that updates are sparse - we only update specific positions. Let me think about this differently. I'll use two kernels: 1. Copy kernel: Copy cache_k and cache_v to updated_k and updated_v 2. RoPE update kernel: For each (b, t, h), apply RoPE rotation and store Actually, a cleaner approach: since the clone() is just element-wise copy, I can fuse it. But the challenge is that the positions can overlap (e.g., multiple t values for the same batch could map to the same position if there's wraparound). Let me check the reference more carefully... Looking at the reference: ```python for b in range(k_new.shape[0]): for t in range(k_new.shape[1]): pos = int(positions[b, t].item()) ... updated_k[b, pos] = rotated updated_v[b, pos] = v_new[b, t] ``` So for each (b, t), we update position positions[b, t]. This is sequential in the reference. If multiple t's map to the same pos for the same batch, later t's would overwrite earlier ones. For parallelization, I need to be careful. Since the number of updates is update_len (16) per batch, and positions could potentially collide, I should handle this properly. Actually, looking at the test setup: ```python base = torch.arange(update_len, dtype=torch.int32).unsqueeze(0).repeat(batch_size, 1) positions = (base + torch.tensor([[7], [31]], dtype=torch.int32)) % cache_len ``` So positions[b, t] = (t + offset[b]) % cache_len. For batch 0, positions are 7, 8, 9, ..., 22. For batch 1, positions are 31, 32, 33, ..., 46. These don't collide within a batch. So I can safely parallelize over (b, t, h) for the RoPE updates, and separately copy the cache. Let me write the full solution: ```python import torch from torch.utils.cpp_extension import load_inline cuda_source = """ #include <torch/extension.h> #include <cuda.h> #include <cuda_runtime.h> template <typename scalar_t> __global__ void copy_cache_kernel( scalar_t* __restrict__ dst, const scalar_t* __restrict__ src, int total_elements ) { int idx = blockIdx.x * blockDim.x + threadIdx.x; if (idx >= total_elements) return; dst[idx] = src[idx]; } template <typename scalar_t> __global__ void rope_update_kernel( scalar_t* __restrict__ updated_k, scalar_t* __restrict__ updated_v, const scalar_t* __restrict__ k_new, const scalar_t* __restrict__ v_new, const scalar_t* __restrict__ cos, const scalar_t* __restrict__ sin, const int* __restrict__ positions, int batch_size, int update_len, int num_heads, int head_dim, int cache_len ) { // Each thread handles one (b, t, h) combination int idx = blockIdx.x * blockDim.x + threadIdx.x; int total_threads = batch_size * update_len * num_heads; if (idx >= total_threads) return; int h = idx % num_heads; idx /= num_heads; int t = idx % update_len; int b = idx / update_len; int pos = positions[b * update_len + t]; // Apply RoPE rotation to key int half_dim = head_dim / 2; int new_k_offset = ((b * cache_len + pos) * num_heads + h) * head_dim; int input_k_offset = ((b * update_len + t) * num_heads + h) * head_dim; int cos_offset = pos * half_dim; for (int d = 0; d < half_dim; d++) { scalar_t c = cos[cos_offset + d]; scalar_t s = sin[cos_offset + d]; scalar_t x_even = k_new[input_k_offset + d * 2]; scalar_t x_odd = k_new[input_k_offset + d * 2 + 1]; updated_k[new_k_offset + d * 2] = x_even * c - x_odd * s; updated_k[new_k_offset + d * 2 + 1] = x_even * s + x_odd * c; } // Copy new value int new_v_offset = ((b * cache_len + pos) * num_heads + h) * head_dim; int input_v_offset = ((b * update_len + t) * num_heads + h) * head_dim; for (int d = 0; d < head_dim; d++) { updated_v[new_v_offset + d] = v_new[input_v_offset + d]; } } std::vector<torch::Tensor> rope_update_cuda( torch::Tensor k_new, torch::Tensor v_new, torch::Tensor cos, torch::Tensor sin, torch::Tensor cache_k, torch::Tensor cache_v, torch::Tensor positions ) { int batch_size = k_new.size(0); int update_len = k_new.size(1); int num_heads = k_new.size(2); int head_dim = k_new.size(3); int cache_len = cache_k.size(1); // Create output tensors auto updated_k = torch::empty_like(cache_k); auto updated_v = torch::empty_like(cache_v); int cache_size = batch_size * cache_len * num_heads * head_dim; const int threads_per_block = 256; // Copy cache int copy_blocks = (cache_size + threads_per_block - 1) / threads_per_block; AT_DISPATCH_FLOATING_TYPES(k_new.scalar_type(), "copy_cache_kernel", ([&] { copy_cache_kernel<scalar_t><<<copy_blocks, threads_per_block>>>( updated_k.data_ptr<scalar_t>(), cache_k.data_ptr<scalar_t>(), cache_size ); copy_cache_kernel<scalar_t><<<copy_blocks, threads_per_block>>>( updated_v.data_ptr<scalar_t>(), cache_v.data_ptr<scalar_t>(), cache_size ); })); // Apply RoPE updates int update_threads = batch_size * update_len * num_heads; int update_blocks = (update_threads + threads_per_block - 1) / threads_per_block; AT_DISPATCH_FLOATING_TYPES(k_new.scalar_type(), "rope_update_kernel", ([&] { rope_update_kernel<scalar_t><<<update_blocks, threads_per_block>>>( updated_k.data_ptr<scalar_t>(), updated_v.data_ptr<scalar_t>(), k_new.data_ptr<scalar_t>(), v_new.data_ptr<scalar_t>(), cos.data_ptr<scalar_t>(), sin.data_ptr<scalar_t>(), positions.data_ptr<int>(), batch_size, update_len, num_heads, head_dim, cache_len ); })); return {updated_k, updated_v}; } """ cpp_source = """ std::vector<torch::Tensor> rope_update_cuda( torch::Tensor k_new, torch::Tensor v_new, torch::Tensor cos, torch::Tensor sin, torch::Tensor cache_k, torch::Tensor cache_v, torch::Tensor positions ); """ rope_extension = load_inline( name='rope_extension', cpp_sources=cpp_source, cuda_sources=cuda_source, functions=['rope_update_cuda'], extra_cuda_cflags=['-O3'], verbose=True ) class ModelNew(nn.Module): def __init__(self): super().__init__() def forward(self, k_new, v_new, cos, sin, cache_k, cache_v, positions): updated_k, updated_v = rope_extension.rope_update_cuda( k_new, v_new, cos, sin, cache_k, cache_v, positions ) return torch.stack((updated_k, updated_v), dim=0) ``` Let me refine this. I want to be more efficient. The copy kernel can be optimized, and I should check if the positions are unique within each batch to ensure correctness. Actually, I realize I should also handle the case where multiple threads might write to the same position if there are collisions. But based on the problem setup, this shouldn't happen. Let me just write a clean implementation first and optimize later.

Good, the kernel compiled. Now let me check correctness.

There's a type issue with the positions tensor. Let me check - positions is expected to be int32, but I'm reading it as int in the kernel. Let me fix this and also make sure the type casting is correct.

Good, the kernel compiled. Let me check correctness again.

The kernel is working correctly. Now let me optimize it further. The current approach has a couple of inefficiencies: 1. Two separate copy kernels for cache_k and cache_v - could merge into one 2. The loop over `d` in the thread - could vectorize with wider loads 3. Could use shared memory to cache cos/sin values Let me create an optimized version that: 1. Uses vectorized loads/stores 2. Uses shared memory for cos/sin when possible 3. Fuses operations better 4. Potentially uses cooperative groups for better memory access patterns Actually, looking at the problem again: - We need to copy the entire cache (which can be large) - We need to update only a small portion (update_len positions) For the copy, cudaMemcpyAsync would be optimal but we're doing CUDA kernels. Using larger vector loads (float4) would help. Let me optimize the kernel:

Let me run correctness to make sure this optimized version works.

Good, correctness passes. Let me try to optimize further by combining both copy and update into a single kernel and reducing memory operations. Since most of the cache is copied unchanged, I could potentially avoid copying altogether by using the same memory and only updating the necessary positions. But the reference creates clones, so I need to maintain that behavior. Actually, wait - the reference calls `clone()` which creates copies. But the returned output is a stack of updated_k and updated_v. If we avoid cloning and directly modify the cache, that might cause issues with gradients and PyTorch's autograd. Let me think about this differently. The key insight is that we only need to update a few positions in the cache (update_len positions per batch). So instead of copying the entire cache, we could: 1. For output positions that don't need updating, copy from input cache 2. For output positions that need updating, apply RoPE and write But this requires knowing which positions need updating. Since positions can be scattered, we'd need to either: - Use atomic operations to mark which positions are updated - Loop through all cache positions and check if they need updating Actually, since update_len << cache_len (16 vs 128), copying the entire cache is wasteful. Let me try a different approach where we: 1. Initialize output with clone (this might be hard to avoid due to PyTorch semantics) 2. Only apply updates But PyTorch's `clone()` should be pretty efficient. Let me try using cudaMemcpyAsync via cudaMemcpy for better throughput, and also try to batch the cache copies into one operation. Actually, let me reconsider. The simplest optimization that should give good results is: 1. Use cudaMemcpy3D for efficient 4D tensor copy 2. Combine the two cache copies into one kernel that does both 3. Increase parallelism in the RoPE kernel Let me try a kernel where we copy both caches and apply RoPE in a single launch, using shared memory to cache cos/sin values.