reasoning
We need to produce a ModelNew subclass or new class that implements same forward but with custom CUDA kernel using torch.utils.cpp_extension.load_inline.
First, understand the operation: For each batch b, each time step t in update_len, get position pos = positions[b, t] (int). Then rotate k_new[b, t] using cos[pos], sin[pos]. The rotation dimension is head_dim. Since head_dim is even, we pair up even and odd indices: x_even = x[..., 0::2]; x_odd = x[..., 1::2]; rot_even = x_even * c - x_odd * s; rot_odd = x_even * s + x_odd * c. Then assign to updated_k[b, pos] = rotated; updated_v[b, pos] = v_new[b, t]. The output is a stack of updated_k, updated_v along dimension 0: shape [2, B, cache_len, H, D].
This is essentially an in-place update of caches with rotary embedding on the new keys.
The naive PyTorch implementation uses Python loops over batch and time, which is very slow. We can write a fused CUDA kernel that does all of this in one pass.
We need to make it faster than the reference. The reference is sequential loops in Python; any parallel kernel will be much faster. But to maximize speedup, we need to write an efficient kernel. Since the dimensions are small (batch=2, update_len=16, num_heads=4, head_dim=32, cache_len=128), total elements are small. However, we can still write a straightforward kernel. The limiting factor might be kernel launch overhead. So we want to keep it simple and coalesced.
Let me think about the kernel design:
Grid: We can launch one thread per element or per pair. Since head_dim=32, we can assign each thread to one element of the update. The total number of elements to update is:
- For k: B * update_len * num_heads * head_dim
- For v: B * update_len * num_heads * head_dim
Total 2 * B * update_len * num_heads * head_dim. That's 2 * 2 * 16 * 4 * 32 = 8192 elements. Very small. But the reference is also small. However, the reference loops in Python, which is extremely slow due to interpreter overhead. So even a naive kernel will be orders of magnitude faster.
But still, we want to make it good. We can write a single kernel that handles both k and v updates, or two kernels. Since the output is stacked, we might write one kernel that writes to two output tensors, or writes to a single tensor of shape [2, B, cache_len, H, D]. It might be easier to allocate the output tensor of shape [2, B, cache_len, H, D] and then launch a kernel that writes into it. Or we can clone cache_k and cache_v into updated_k and updated_v and then launch a kernel that overwrites the relevant positions.
Wait, we need to return both updated caches, where all previous entries are kept and the new positions are overwritten. So we can do:
updated_k = cache_k.clone()
updated_v = cache_v.clone()
Then call kernel that updates updated_k and updated_v at positions[b, t] with rotated k_new and v_new.
Alternatively, we can allocate output = torch.empty(2, B, cache_len, H, D), copy cache_k/cache_v into output[0], output[1], and then update. But clone is fine.
We need to write a CUDA kernel using `load_inline`. The kernel signature can be something like:
void rope_kv_update(const float* k_new, const float* v_new, const float* cos, const float* sin, const int* positions, float* updated_k, float* updated_v, int B, int T, int cache_len, int H, int D)
Where T = update_len.
In the kernel, each thread can handle one (b, t, h) pair. For each such pair, we read pos = positions[b * T + t], then for each d in 0..D-1, we compute rotated_k and write to updated_k[b * cache_len * H * D + pos * H * D + h * D + d] = ..., and write v_new directly to updated_v[...].
But D is small (32). So we can have each thread handle one (b, t, h, d) element. Let's do that. Total threads = B * T * H * D. For k, each thread computes one element. For even d, it needs x_even and x_odd. That means if each thread handles one d, for even d it needs to read two elements from k_new to compute the rotation. That's fine but requires reading both x_even and x_odd and the corresponding c and s. But c and s are per position and per pair index. Let's define pair_idx = d / 2. c = cos[pos * (D/2) + pair_idx]; s = sin[...].
If d is even:
x_even = k_new[b,t,h,d]
x_odd = k_new[b,t,h,d+1]
rot = x_even * c - x_odd * s
If d is odd:
x_even = k_new[b,t,h,d-1]
x_odd = k_new[b,t,h,d]
rot = x_even * s + x_odd * c
This reading is efficient if threads in a warp access adjacent d's, because for even d, thread d reads d and d+1; for odd d, thread d reads d-1 and d. This causes some bank conflicts or uncoalesced reads? Let's think: k_new is stored in row-major order, so elements for a given (b,t,h) are contiguous in d. If we have a warp processing d from 0 to 31, thread 0 reads d=0 and d=1, thread 1 reads d=1 and d=0? Wait d=1 is odd, so reads d=0 and d=1. So each element is read twice. That's okay for small D. But could also just have each thread handle a pair (process two elements). Let's have each thread handle a pair: thread idx handles d_pair = 0..D/2 - 1. Then each thread computes both even and odd outputs and writes them. Number of threads per (b,t,h) is D/2 = 16. Total threads = B * T * H * (D/2) = 2 * 16 * 4 * 16 = 2048. This is even smaller. But warp size is 32, so we can have each warp handle two (b,t,h) blocks. This might be fine.
Alternatively, we can have each thread handle one (b,t,h) and loop over D/2 pairs. This might be better because it reduces total thread count but increases work per thread. Since total work is small, having more threads might be okay but could lead to underutilization. However, any approach will be fast because the reference is Python loops. Let's think about the typical optimization for small workloads: minimize kernel launch overhead. A single kernel with a grid of (B, T, H) blocks and each block with D/2 threads = 16 threads. That's tiny blocks. Or we can use a 1D grid and block size 256, with each thread handling multiple items.
Simplify: use a 1D kernel where each thread handles one (b, t, h, d_pair). This is straightforward.
Total pair count = B * T * H * (D/2). For our typical sizes, that's 2048. If we use block size 256, we need 8 blocks. That launches 8 blocks of 256 threads = 2048 threads. Each block gets one chunk of pairs. This is very reasonable and coalesced.
Compute index:
int pair_idx = threadIdx.x + blockIdx.x * blockDim.x;
int total_pairs = B * T * H * (D/2);
if (pair_idx >= total_pairs) return;
Then decode:
int half_D = D / 2;
int d = pair_idx % half_D;
int tmp = pair_idx / half_D;
int h = tmp % H;
tmp /= H;
int t = tmp % T;
int b = tmp / T;
Then pos = positions[b * T + t];
Then read x_even = k_new[((b * T + t) * H + h) * D + 2*d];
x_odd = k_new[... + 2*d + 1];
c = cos[pos * half_D + d];
s = sin[pos * half_D + d];
rot_even = x_even * c - x_odd * s;
rot_odd = x_even * s + x_odd * c;
Write to updated_k:
int k_base = ((b * cache_len + pos) * H + h) * D;
updated_k[k_base + 2*d] = rot_even;
updated_k[k_base + 2*d + 1] = rot_odd;
And for v:
int v_base = ((b * cache_len + pos) * H + h) * D;
updated_v[v_base + 2*d] = v_new[((b * T + t) * H + h) * D + 2*d];
updated_v[v_base + 2*d + 1] = v_new[... + 2*d + 1];
Wait, v_new is just copied as-is. So we can also have each thread copy two floats. That's fine.
But we also need to handle the cloning. In forward, we can do:
updated_k = cache_k.clone()
updated_v = cache_v.clone()
Then call the kernel to overwrite the relevant positions.
Alternatively, we can avoid clone by allocating output and using a separate kernel to copy the old cache, then write the new data. But clone is a fast memcpy, and cache_len=128, so it's tiny. Using clone is fine.
One nuance: positions is int32. We should accept it as int32 (or int64, but the get_inputs uses torch.int32). In CUDA, we can read it as int.
Also note, the cos and sin tables are shape [cache_len, head_dim//2]. In get_inputs, they are generated as float32.
We need to make sure we write the kernel and host wrapper correctly.
Given that we're using `load_inline`, the CUDA source should be in `cuda_sources`. The host wrapper in `cpp_sources`. `functions` should list the wrapper function.
Let's draft the wrapper:
void rope_kv_update(
const torch::Tensor& k_new,
const torch::Tensor& v_new,
const torch::Tensor& cos,
const torch::Tensor& sin,
const torch::Tensor& positions,
torch::Tensor& updated_k,
torch::Tensor& updated_v) {
int B = k_new.size(0);
int T = k_new.size(1);
int H = k_new.size(2);
int D = k_new.size(3);
int cache_len = updated_k.size(1);
int half_D = D / 2;
int total_pairs = B * T * H * half_D;
const int threads = 256;
const int blocks = (total_pairs + threads - 1) / threads;
rope_kv_update_kernel<<<blocks, threads>>>(
k_new.data_ptr<float>(),
v_new.data_ptr<float>(),
cos.data_ptr<float>(),
sin.data_ptr<float>(),
positions.data_ptr<int>(), // or int32_t
updated_k.data_ptr<float>(),
updated_v.data_ptr<float>(),
B, T, cache_len, H, D);
}
Need to include <torch/extension.h> or <ATen/ATen.h>. Since we are using load_inline, we can use <torch/extension.h>.
Make sure to qualify namespaces: at::Tensor, torch::Tensor.
Wait, `cpp_sources` should contain the host function definitions. `cuda_sources` should contain the CUDA kernel. But in `load_inline`, `cpp_sources` is typically for C++ host code and `cuda_sources` for CUDA code. Since the host wrapper will call the CUDA kernel, it needs to declare the kernel. We can put the kernel declaration in `cpp_sources` and definition in `cuda_sources`. Or we can put everything in `cuda_sources`? `load_inline` supports both. Usually `cpp_sources` is for C++ code and `cuda_sources` for CUDA code. The host function can be in `cpp_sources` but it needs the kernel declaration. We can provide the kernel prototype in `cpp_sources` and the definition in `cuda_sources`. But it's also okay to put the host wrapper in `cpp_sources` and the kernel in `cuda_sources`, and the compiler will link them because they are compiled together.
Actually, `torch.utils.cpp_extension.load_inline` compiles all sources together. If we define the kernel in `cuda_sources` and call it from `cpp_sources`, the compiler must see the kernel declaration. In NVCC, device functions are visible across translation units? Not necessarily, but `load_inline` might combine them or compile `cpp_sources` with the host compiler and `cuda_sources` with nvcc, then link. If we put the host function in `cpp_sources`, the C++ compiler won't understand `<<< >>>` syntax. Wait, `cpp_sources` is compiled as C++ code (not .cu) by the host compiler if it doesn't have .cu extension? Actually `load_inline` takes `cpp_sources` which is compiled as C++ source files. But calling CUDA kernels from C++ requires the file to be compiled as CUDA (i.e., .cu). However, `cpp_sources` is typically compiled with the C++ compiler, so you cannot use `<<< >>>` inside it. In `load_inline`, sometimes you can include CUDA code in `cuda_sources`, and the host wrapper can be in `cpp_sources` if you use `AT_CUDA_CHECK` and a helper template or if you write a wrapper in `cpp_sources` that calls a CUDA function. But actually, PyTorch's `load_inline` is a bit flexible: both `cpp_sources` and `cuda_sources` are written to temporary files with `.cpp` and `.cu` extensions, respectively, and compiled with the appropriate compiler. You cannot use `<<< >>>` in `.cpp` files because the host compiler doesn't understand it. But you can declare the kernel as `void my_kernel(...);` and then call it using a launcher function defined in the `.cu` file, or use the `AT_DISPATCH` pattern.
Alternatively, you can put the host wrapper in the `cuda_sources` as well? The `cuda_sources` is a list of strings, each compiled as `.cu`. You can put host-side C++ code in `.cu` files; nvcc handles both host and device code. So you could put everything in `cuda_sources` and leave `cpp_sources` empty. But `functions=[...]` needs to list the host functions defined in either source. Actually, the docs say: `functions` is a list of function names to bind with pybind11. These functions can be defined in `cpp_sources` or `cuda_sources`.
So the simplest approach: put everything (kernel and host wrapper) in `cuda_sources`. Then `cpp_sources` can be empty or omitted. Wait, `load_inline` signature: `load_inline(name, cpp_sources, cuda_sources=None, functions=None, ...)`. If we provide only `cuda_sources`, it's fine. nvcc will compile the `.cu` file containing both the kernel and the host function. The host function is compiled as host code, and pybind11 will bind it. This is the easiest.
Let's verify: The doc says: "cpp_sources (str or list) – C++ sources. cuda_sources (str or list) – CUDA sources." But NVCC compiles `.cu` files. If you put C++ host code with pybind11 bindings inside a `.cu` file, nvcc will compile it. That works. Actually, it's common to put everything in `cpp_sources` and use `AT_DISPATCH_FLOATING_TYPES_AND_HALF` macros, but that's because the kernel is defined with `__global__` inside a macro. No, that's for `cpp_extension.load` where the file is `.cpp` but contains CUDA code? Hmm. Let's be safe.
We can write `cuda_sources` containing:
#include <torch/extension.h>
#include <cuda_runtime.h>
__global__ void rope_kv_update_kernel(...)
void rope_kv_update(const torch::Tensor& k_new, ...) {
...
rope_kv_update_kernel<<<blocks, threads>>>(...);
}
Then in `functions=['rope_kv_update']`.
But the prompt says: "Do not write your own `PYBIND11_MODULE(...)` block in `cpp_sources`. `load_inline` auto-generates one from the `functions=[...]` argument; including your own causes a duplicate-symbol redefinition error at compile time."
So we must not write `PYBIND11_MODULE`. We just define the host function and list it in `functions`.
Now, should we put the host function in `cpp_sources` or `cuda_sources`? Since it calls a CUDA kernel, it's easiest to put it in `cuda_sources`. Because if we put it in `cpp_sources`, we'd need to separate the kernel launch into an `extern "C"` helper or something. Actually, `cpp_sources` can call a helper function declared in `cuda_sources`. But again, simpler to put everything in `cuda_sources`.
But wait: `load_inline` generates a `PYBIND11_MODULE` that binds the functions listed. If we define the function in `cuda_sources`, that's fine. The function will have C++ linkage (not `extern "C"` needed because pybind11 uses mangled names via templates? Actually pybind11's `m.def("rope_kv_update", &rope_kv_update, ...)` uses the C++ name and gets the function pointer. If the function is defined in a `.cu` file compiled by nvcc, its mangled name is standard C++. That should link fine with the `PYBIND11_MODULE` generated in the `.cpp` file. But there could be issues with linking across host compiler and nvcc name mangling if they differ (e.g., gcc vs nvcc with different ABI). Usually they are compatible. To be safe, we can declare the host function as `extern "C" void rope_kv_update(...)` in the `.cu` file. But pybind11 needs the C++ type to bind it. `extern "C"` gives C linkage but the function still has C++ types. pybind11 can handle that because it uses the function pointer, not the name. However, it uses template deduction based on the function type, and the symbol name doesn't matter. So `extern "C"` is fine, but not necessary.
I think putting everything in `cuda_sources` is fine.
Let's write the kernel.
Inputs:
- k_new: float*, shape [B, T, H, D]
- v_new: float*, shape [B, T, H, D]
- cos: float*, shape [cache_len, half_D]
- sin: float*, shape [cache_len, half_D]
- positions: int*, shape [B, T]
- updated_k: float*, shape [B, cache_len, H, D]
- updated_v: float*, shape [B, cache_len, H, D]
Parameters: B, T, cache_len, H, D.
Kernel:
__global__ void rope_kv_update_kernel(
const float* __restrict__ k_new,
const float* __restrict__ v_new,
const float* __restrict__ cos,
const float* __restrict__ sin,
const int* __restrict__ positions,
float* __restrict__ updated_k,
float* __restrict__ updated_v,
int B,
int T,
int cache_len,
int H,
int D)
{
int half_D = D >> 1;
int total_pairs = B * T * H * half_D;
int idx = blockIdx.x * blockDim.x + threadIdx.x;
if (idx >= total_pairs) return;
// Decode indices
int d = idx % half_D;
int tmp = idx / half_D;
int h = tmp % H;
tmp /= H;
int t = tmp % T;
int b = tmp / T;
int pos = positions[b * T + t];
// read k_new
int k_new_stride_b = T * H * D;
int k_new_stride_t = H * D;
int k_new_stride_h = D;
int k_new_idx = b * k_new_stride_b + t * k_new_stride_t + h * k_new_stride_h;
float x_even = k_new[k_new_idx + 2 * d];
float x_odd = k_new[k_new_idx + 2 * d + 1];
int table_idx = pos * half_D + d;
float c = cos[table_idx];
float s = sin[table_idx];
float rot_even = x_even * c - x_odd * s;
float rot_odd = x_even * s + x_odd * c;
int cache_stride_b = cache_len * H * D;
int cache_stride_pos = H * D;
int cache_stride_h = D;
int cache_idx = b * cache_stride_b + pos * cache_stride_pos + h * cache_stride_h;
updated_k[cache_idx + 2 * d] = rot_even;
updated_k[cache_idx + 2 * d + 1] = rot_odd;
updated_v[cache_idx + 2 * d] = v_new[k_new_idx + 2 * d];
updated_v[cache_idx + 2 * d + 1] = v_new[k_new_idx + 2 * d + 1];
}
But maybe we should use `int64_t` for indexing to avoid overflow? For small dimensions, int is fine. But it's good practice to use `int64_t` when needed. Here max index is B*cache_len*H*D = 2*128*4*32 = 32768. So int is fine.
We need to ensure contiguous memory? The input tensors might not be contiguous. It's safer to call `.contiguous()` before passing to the kernel. Or we can use TensorAccessor. But `.contiguous()` is fine and cheap for these sizes.
Also, we need to check if positions is int32 or int64. In get_inputs, positions is torch.int32. So `positions.data_ptr<int>()` is fine. But if someone passes int64, it'll break. We can cast to int: `positions.to(torch::kInt32)`. Or we can read as int32_t. Let's ensure in the wrapper: `positions = positions.to(torch::kInt32);`. Since the reference uses `.item()` which handles int32, it's fine.
But wait, the reference gets positions as `torch.int32`. So we can assume that. However, to be robust, we can do `TORCH_CHECK(positions.dtype() == torch::kInt32, "positions must be int32");` or convert. Converting is safer.
Similarly, we should ensure all tensors are on CUDA? The prompt doesn't specify device. In the reference, inputs are CPU tensors (since get_inputs doesn't put them on GPU). But the evaluation will likely run on GPU? Actually, the model reference is a nn.Module; we need to produce a ModelNew that can run on CPU or GPU? The custom CUDA kernel requires CUDA tensors. The evaluation harness might move inputs to GPU. Or it might test on CPU? But they ask for a custom CUDA kernel, so presumably inputs will be on CUDA device during timing. However, the reference implementation uses pure PyTorch operations that work on both CPU and GPU. For our ModelNew, if inputs are on CPU, the kernel will fail. But timing is on GPU. To be safe, we should check device and if it's CPU, maybe fall back to reference? Wait, the instruction says "Do not use try/except fallbacks to the reference". But we can check device and if not CUDA, run a PyTorch equivalent. Is that considered a fallback? It's not try/except. But static check might catch it if it looks like fallback. The prompt says: "Do not use try/except fallbacks to the reference, ...". We can have an if statement: if not k_new.is_cuda(): return ... (reference logic). But that's a fallback. The static check might flag it. Let's avoid that. Instead, we can just run the kernel and assume inputs are on CUDA. The evaluation harness will likely call `.cuda()` before timing. In the reference, `get_inputs()` returns CPU tensors, but the harness may move them.
Wait, the reference uses `.item()` inside loops, which forces CPU sync. That's part of why it's slow. If the inputs were on GPU, `.item()` would be even slower. The harness will likely run the reference on CPU? But the task says "custom CUDA kernel that replaces a PyTorch reference... faster than the reference." Usually they test on GPU. The reference is CPU-based. So we'll assume inputs are on GPU. Or maybe they run reference on GPU too? On GPU, the Python for loop with `.item()` is brutally slow. Our CUDA kernel will be much faster.
But wait, if inputs are on CPU, our custom CUDA kernel cannot run. In that case, we could copy to CUDA, run, copy back. That adds overhead but still might be faster? Actually for tiny sizes, copying to CUDA will dominate. But the evaluation for correctness likely runs on CPU or GPU? In these tasks, they usually test correctness on GPU because the kernel is CUDA. The `run_correctness` and `submit_kernel` will move inputs to the GPU if needed? Let's assume they set the default CUDA device. The reference will run on GPU if inputs are on GPU, but the Python loops will be on CPU (since `.item()` brings to CPU). Actually `positions[b,t].item()` on a GPU tensor does a device-to-host copy for each element, which is extremely slow. But the reference implementation is what it is. Our kernel should run on GPU.
I think it's safe to assume inputs are CUDA tensors during evaluation. Let's double-check the reference: `k_new = torch.randn(...)` no `.cuda()`. But the harness might call `k_new = k_new.cuda()` before passing. Usually in these tasks, the inputs are on GPU. The `get_inputs` returns CPU tensors, but the evaluation code does something like:
inputs = get_inputs()
if torch.cuda.is_available():
inputs = [x.cuda() for x in inputs]
model = Model().cuda()
But the reference Model doesn't have parameters, so `.cuda()` on the module does nothing; but the inputs would be on GPU. Then inside the `forward`, `positions[b,t].item()` would copy each scalar from GPU to CPU. That's the bottleneck.
Okay, so we need to write a ModelNew that clones the caches on the same device as inputs, and calls the CUDA kernel.
One concern: In the reference, `updated_k = cache_k.clone()` means if cache_k is on GPU, updated_k is on GPU. Good.
So our `forward` can do:
updated_k = cache_k.clone()
updated_v = cache_v.clone()
# Ensure contiguous
k_new_c = k_new.contiguous()
v_new_c = v_new.contiguous()
cos_c = cos.contiguous()
sin_c = sin.contiguous()
pos_c = positions.to(torch.int32).contiguous()
# call kernel
_ = torch.cuda.Stream() maybe not needed.
module.rope_kv_update(k_new_c, v_new_c, cos_c, sin_c, pos_c, updated_k, updated_v)
return torch.stack((updated_k, updated_v), dim=0)
Wait, the kernel overwrites updated_k and updated_v in-place. The host wrapper takes torch::Tensor references. In Python, the extension function can accept Tensor objects. But if the wrapper signature is `void rope_kv_update(const torch::Tensor& k_new, ..., torch::Tensor& updated_k, ...)`, then in Python, we need to pass all arguments. The pybind11 wrapping should handle non-const references for output tensors.
However, pybind11 by default passes arguments by value for `torch::Tensor`? Actually, `torch::Tensor` is a reference-counted pointer. Passing by const reference or by value is fine. But to get an output parameter that is modified in place, we can pass it by non-const reference. In pybind11, `m.def("rope_kv_update", &rope_kv_update, ...)` with arguments (const Tensor&, ..., Tensor&, Tensor&) will allow Python to pass the tensor and modifications will be reflected because `torch::Tensor` is a handle to the underlying storage. However, pybind11 might not allow binding functions with non-const references unless specified? Actually, it does: it passes the Python object by reference, and the C++ function gets a reference to the C++ object. For `torch::Tensor`, it's copyable; but if we take `Tensor&`, pybind11 will pass a reference to the Tensor object held by the Python argument. Since Tensor is a reference type, modifying the data it points to is fine. But we must be careful: pybind11 might require `py::arg().noconvert()` or something? Actually no, for `Tensor&` it's fine: it binds as a reference to the existing object.
Alternatively, we can have the function return a tuple of (updated_k, updated_v) after copying. But we already cloned them; in-place modification is fine.
But wait: `torch::Tensor` is a shared pointer to `TensorImpl`. If we pass `Tensor&`, the function modifies the data of the tensor, not the tensor metadata. Python sees the data change. That's okay.
We could also define the function to return void and modify in place.
Alternatively, to avoid potential pybind11 issues with non-const reference of a pybind11-bound type, we can pass the tensors by value: `torch::Tensor updated_k`. Wait, passing by value would copy the Tensor handle, but modifications to the data would still affect the same underlying storage. So `void rope_kv_update(..., torch::Tensor updated_k, torch::Tensor updated_v)` is fine and safer for pybind11. Let's do that.
Wait, but pybind11 might treat `torch::Tensor` as an opaque type and pass by reference anyway. In PyTorch's binding, functions often take `const at::Tensor&` for inputs and `at::Tensor` for outputs. I think `at::Tensor` by value is okay. Let's do:
void rope_kv_update(const at::Tensor& k_new,
const at::Tensor& v_new,
const at::Tensor& cos,
const at::Tensor& sin,
const at::Tensor& positions,
at::Tensor updated_k,
at::Tensor updated_v) {
...
}
This should work.
We also need to handle the case where `k_new` etc may not be contiguous. We can call `.contiguous()` in Python before passing. So the kernel can assume contiguous.
Now, the kernel should check that all tensors are on CUDA device? Not necessary if we assume.
Let's also consider using `__restrict__` to help compiler.
Given the small problem size, we might also consider using a single block with shared memory? But unnecessary.
Alternatively, we could write the kernel to process each (b, t, h) with a thread and loop over d. This might have less thread divergence and better instruction-level parallelism. Let's consider:
Each thread does:
int idx = blockIdx.x * blockDim.x + threadIdx.x;
int total = B * T * H;
if (idx >= total) return;
decode b, t, h.
pos = positions[b*T + t];
half_D = D/2;
k_new_idx = ((b * T + t) * H + h) * D;
cache_idx = ((b * cache_len + pos) * H + h) * D;
table_base = pos * half_D;
for (int d = 0; d < half_D; ++d) {
float x_even = k_new[k_new_idx + 2*d];
float x_odd = k_new[k_new_idx + 2*d + 1];
float c = cos[table_base + d];
float s = sin[table_base + d];
updated_k[cache_idx + 2*d] = x_even * c - x_odd * s;
updated_k[cache_idx + 2*d + 1] = x_even * s + x_odd * c;
updated_v[cache_idx + 2*d] = v_new[k_new_idx + 2*d];
updated_v[cache_idx + 2*d + 1] = v_new[k_new_idx + 2*d + 1];
}
Total threads = B*T*H = 2*16*4 = 128. That's very small. One block of 128 threads. But loops over 16 pairs. This might be fine and even better because each thread reads positions once and does sequential memory access. The memory per thread is 32 floats for k_new, 32 for v_new, 32 for updated_k, 32 for updated_v. That's 128*128*4 = 64KB of reads/writes. Within a warp, adjacent threads access different (b,t,h). Memory access for k_new: thread 0 reads b=0,t=0,h=0; thread 1 reads b=0,t=0,h=1. If H=4, threads 0-3 read h=0..3 for same b,t. The k_new data for these h are contiguous: k_new[b,t,0,:], k_new[b,t,1,:], etc. So a warp of 32 threads reads 8 different (b,t) blocks (since H=4, 32 threads cover 32/4 = 8 t's). The memory pattern is strided by H*D=128 floats. So within a warp, each thread reads a different 128-float chunk. That's not coalesced. But because each chunk is 128*4=512 bytes, and there are 32 threads, the total per warp is 16KB. The memory access is scattered. However, the L1 cache might help. But with only 128 threads, coalescing is less critical.
Alternatively, if we use 2048 threads (one per pair), thread 0 reads pair 0 of h=0,t=0,b=0. Thread 1 reads pair 1 of same h,t,b. They are adjacent in memory: k_new_idx + 2*d. So threads in a warp read adjacent floats -> coalesced. That's better.
The pair-based approach seems better for memory coalescing. Let's stick with that.
Total threads = 2048. Block size 256. 8 blocks.
Memory accesses:
- k_new: for a warp processing 32 consecutive pair indices, what are the addresses?
pair_idx = blockIdx.x*256 + warp_id*32 + lane.
Suppose warp lane 0..31. If half_D=16, then within a warp, pairs belong to the same (b,t,h) if all 32 lanes are within same half_D. Since half_D=16, a warp covers two consecutive (b,t,h) groups: lanes 0..15 belong to (b,t,h) group A, d=0..15; lanes 16..31 belong to next group (b,t,h)+1, d=0..15.
So thread lane 0 reads address of pair 0 of group A. Lane 1 reads pair 1 of group A. These are contiguous floats (2*4=8 bytes apart). That's coalesced (each thread reads 2 floats? Actually each thread reads x_even and x_odd, which are at offsets 2d and 2d+1). So thread lane d reads floats at offsets 2d and 2d+1. That's contiguous. The warp issues 16 loads of 2 contiguous floats? Actually each thread loads two floats. The addresses across the warp are [0,1], [2,3], ..., [30,31]. That's perfectly coalesced into two 128-byte transactions (or one 128-byte for 32 floats?). Actually a warp of 32 threads each loading 4 bytes = 128 bytes. Here each thread loads two 4-byte floats = 8 bytes per thread = 256 bytes per warp. With L1 caching, it might be two 128-byte segments. So coalescing is good.
For the next 16 lanes (16..31), they read the first 32 floats of the next (b,t,h) group. Those addresses are offset by H*D = 128 floats = 512 bytes from group A. So the warp's accesses are split into two segments: one at offset X, size 64 floats, and one at offset X+512, size 64 floats. Because a warp executes together, the addresses are interleaved? Actually threads in a warp have lanes 0..31. The memory controller will handle two separate segments. Still fine.
Then writing to updated_k: similar pattern, but pos can be different for different t. So (b,t,h) group A has pos_A, group B (next t or h) has pos_B. The write addresses depend on pos. If pos values are not contiguous across groups, writes may be scattered. But since positions are small, many writes might still hit L2.
Given the tiny workload, this is definitely fine.
We can also consider using vectorized loads/stores (float2) to improve throughput. Let's see if we can load two floats at a time. Each thread naturally reads x_even and x_odd. We can use `float2` to load them in one instruction if the address is 8-byte aligned. Since each pair is at offset 2*d, and d increments by 1 per pair, the start of each pair is at an even float index. For a given (b,t,h), the base address is aligned to the start of the tensor? Not necessarily 8-byte aligned if the base is not a multiple of 8. But `k_new` is a torch tensor, its data pointer is typically aligned to at least 256 bytes or something. So base + 2*d is 8-byte aligned for d even? Actually 2*d * 4 = 8*d. So offset is always a multiple of 8 bytes. So `float2` loads are aligned! That's great.
We can reinterpret the pointer as `const float2*` and load `float2 x = *(const float2*)(k_new + k_new_idx + 2*d);` where `x.x` is even and `x.y` is odd. Then store similarly. For v_new and updated_v, also aligned.
For updated_k, the base address is `cache_idx = b*cache_len*H*D + pos*H*D + h*D`. `cache_idx` offset from the tensor base is a multiple of D? Since D=32, 32*4=128 bytes. So `cache_idx` is 128-byte aligned. Then `cache_idx + 2*d` is 8-byte aligned for any d. So `float2` stores are aligned.
So we can use `float2` to load and store pairs. This halves the number of memory instructions.
Similarly, cos and sin are per pair. We can also load them as `float2`? Cos and sin arrays are separate. We could load cos[pos*half_D + d] and sin[...] individually, or we could pre-interleave cos and sin into a single tensor? But we can just load them as two floats. Since half_D is small, they might be in registers.
Let's write the kernel with float2:
__global__ void rope_kv_update_kernel(
const float* __restrict__ k_new,
const float* __restrict__ v_new,
const float* __restrict__ cos,
const float* __restrict__ sin,
const int* __restrict__ positions,
float* __restrict__ updated_k,
float* __restrict__ updated_v,
int B,
int T,
int cache_len,
int H,
int D)
{
int half_D = D >> 1;
int total_pairs = B * T * H * half_D;
int idx = blockIdx.x * blockDim.x + threadIdx.x;
if (idx >= total_pairs) return;
int d = idx % half_D;
int tmp = idx / half_D;
int h = tmp % H;
tmp /= H;
int t = tmp % T;
int b = tmp / T;
int pos = positions[b * T + t];
int k_new_idx = ((b * T + t) * H + h) * D + 2 * d;
float2 x = reinterpret_cast<const float2*>(k_new)[k_new_idx / 2];
float c = cos[pos * half_D + d];
float s = sin[pos * half_D + d];
float2 rot;
rot.x = x.x * c - x.y * s;
rot.y = x.x * s + x.y * c;
int cache_idx = ((b * cache_len + pos) * H + h) * D + 2 * d;
reinterpret_cast<float2*>(updated_k)[cache_idx / 2] = rot;
reinterpret_cast<float2*>(updated_v)[cache_idx / 2] = reinterpret_cast<const float2*>(v_new)[k_new_idx / 2];
}
Wait, `k_new_idx` is always even because it's `2*d`. So `k_new_idx / 2` is an integer. But if we do integer division, it's fine. Actually we can just compute `int pair_offset = ((b * T + t) * H + h) * half_D + d;` and then load from `reinterpret_cast<const float2*>(k_new)[pair_offset];`. Because the number of pairs per (b,t,h) is half_D. The base pointer is aligned, so indexing by pair number works.
Similarly, `cache_idx` can be computed as pairs: `int cache_pair_offset = ((b * cache_len + pos) * H + h) * half_D + d;`
Then:
const float2* k_new_ptr = reinterpret_cast<const float2*>(k_new);
const float2* v_new_ptr = reinterpret_cast<const float2*>(v_new);
float2* updated_k_ptr = reinterpret_cast<float2*>(updated_k);
float2* updated_v_ptr = reinterpret_cast<float2*>(updated_v);
float2 x = k_new_ptr[pair_offset];
...
updated_k_ptr[cache_pair_offset] = rot;
updated_v_ptr[cache_pair_offset] = v_new_ptr[pair_offset];
This is cleaner and avoids division.
Let's compute pair_offset:
int half_D = D / 2;
int pair_offset = ((b * T + t) * H + h) * half_D + d;
cache_pair_offset = ((b * cache_len + pos) * H + h) * half_D + d;
Yes.
This indexing assumes the tensor is contiguous in the last dimension and all preceding dimensions are standard row-major. Since we will enforce `.contiguous()`, it's fine.
Now, we need to be careful about `reinterpret_cast` and type punning. In CUDA C++, `reinterpret_cast<const float2*>(k_new)` is okay if the pointer is aligned to 8 bytes. As discussed, it is. So this is safe.
Now, what about D not being divisible by 2? The reference uses `0::2` and `1::2`, which implies D is even. For safety, we can assert D%2==0.
Let's write the host wrapper in `cuda_sources`:
#include <torch/extension.h>
#include <cuda_runtime.h>
__global__ void rope_kv_update_kernel(...) { ... }
void rope_kv_update(
const at::Tensor& k_new,
const at::Tensor& v_new,
const at::Tensor& cos,
const at::Tensor& sin,
const at::Tensor& positions,
at::Tensor updated_k,
at::Tensor updated_v) {
// checks
TORCH_CHECK(k_new.is_cuda(), "k_new must be CUDA");
// ... etc.
int B = k_new.size(0);
int T = k_new.size(1);
int H = k_new.size(2);
int D = k_new.size(3);
int cache_len = updated_k.size(1);
int half_D = D / 2;
int total_pairs = B * T * H * half_D;
const int threads = 256;
const int blocks = (total_pairs + threads - 1) / threads;
rope_kv_update_kernel<<<blocks, threads>>>(
k_new.data_ptr<float>(),
v_new.data_ptr<float>(),
cos.data_ptr<float>(),
sin.data_ptr<float>(),
positions.data_ptr<int>(),
updated_k.data_ptr<float>(),
updated_v.data_ptr<float>(),
B, T, cache_len, H, D);
}
Wait, `positions.data_ptr<int>()` assumes dtype is int32. Let's use `int32_t` to be explicit. And we should cast or check: `positions.scalar_type() == at::kInt`. If not, we can convert.
In Python, easier to pass `.to(torch.int32)`.
Also need to ensure D is even. We can assert.
One more thing: The output of forward is `torch.stack((updated_k, updated_v), dim=0)`. This creates a new tensor of shape [2, B, cache_len, H, D].
Now, let's think about using `torch.utils.cpp_extension.load_inline` in module-level code. We need to supply `name`, `cuda_sources`, and `functions`.
Example:
import torch
from torch.utils.cpp_extension import load_inline
cuda_source = r'''
#include <torch/extension.h>
#include <cuda_runtime.h>
__global__ void rope_kv_update_kernel(
const float* __restrict__ k_new,
const float* __restrict__ v_new,
const float* __restrict__ cos,
const float* __restrict__ sin,
const int* __restrict__ positions,
float* __restrict__ updated_k,
float* __restrict__ updated_v,
int B,
int T,
int cache_len,
int H,
int D)
{
int half_D = D >> 1;
int total_pairs = B * T * H * half_D;
int idx = blockIdx.x * blockDim.x + threadIdx.x;
if (idx >= total_pairs) return;
int d = idx % half_D;
int tmp = idx / half_D;
int h = tmp % H;
tmp /= H;
int t = tmp % T;
int b = tmp / T;
int pos = positions[b * T + t];
int pair_offset = ((b * T + t) * H + h) * half_D + d;
int cache_pair_offset = ((b * cache_len + pos) * H + h) * half_D + d;
const float2* k_new_ptr = reinterpret_cast<const float2*>(k_new);
const float2* v_new_ptr = reinterpret_cast<const float2*>(v_new);
const float* cos_ptr = cos;
const float* sin_ptr = sin;
float2* updated_k_ptr = reinterpret_cast<float2*>(updated_k);
float2* updated_v_ptr = reinterpret_cast<float2*>(updated_v);
float2 x = k_new_ptr[pair_offset];
float c = cos_ptr[pos * half_D + d];
float s = sin_ptr[pos * half_D + d];
float2 rot;
rot.x = x.x * c - x.y * s;
rot.y = x.x * s + x.y * c;
updated_k_ptr[cache_pair_offset] = rot;
updated_v_ptr[cache_pair_offset] = v_new_ptr[pair_offset];
}
void rope_kv_update(
const at::Tensor& k_new,
const at::Tensor& v_new,
const at::Tensor& cos,
const at::Tensor& sin,
const at::Tensor& positions,
at::Tensor updated_k,
at::Tensor updated_v)
{
TORCH_CHECK(k_new.is_cuda(), "k_new must be a CUDA tensor");
TORCH_CHECK(v_new.is_cuda(), "v_new must be a CUDA tensor");
TORCH_CHECK(cos.is_cuda(), "cos must be a CUDA tensor");
TORCH_CHECK(sin.is_cuda(), "sin must be a CUDA tensor");
TORCH_CHECK(positions.is_cuda(), "positions must be a CUDA tensor");
TORCH_CHECK(updated_k.is_cuda(), "updated_k must be a CUDA tensor");
TORCH_CHECK(updated_v.is_cuda(), "updated_v must be a CUDA tensor");
int B = k_new.size(0);
int T = k_new.size(1);
int H = k_new.size(2);
int D = k_new.size(3);
int cache_len = updated_k.size(1);
TORCH_CHECK(D % 2 == 0, "head_dim must be even");
int half_D = D / 2;
int total_pairs = B * T * H * half_D;
const int threads = 256;
const int blocks = (total_pairs + threads - 1) / threads;
rope_kv_update_kernel<<<blocks, threads>>>(
k_new.data_ptr<float>(),
v_new.data_ptr<float>(),
cos.data_ptr<float>(),
sin.data_ptr<float>(),
positions.data_ptr<int>(),
updated_k.data_ptr<float>(),
updated_v.data_ptr<float>(),
B, T, cache_len, H, D);
}
'''
module = load_inline(name="rope_kv_cache_cuda", cuda_sources=cuda_source, functions=["rope_kv_update"], extra_cuda_cflags=["-O3"])
Then ModelNew:
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 = cache_k.clone()
updated_v = cache_v.clone()
# Ensure contiguous and correct dtype
k_new_c = k_new.contiguous()
v_new_c = v_new.contiguous()
cos_c = cos.contiguous()
sin_c = sin.contiguous()
pos_c = positions.to(torch.int32).contiguous()
module.rope_kv_update(k_new_c, v_new_c, cos_c, sin_c, pos_c, updated_k, updated_v)
return torch.stack((updated_k, updated_v), dim=0)
Wait, we need to import `module` in the class or make it accessible. Since `load_inline` is at module level, we can assign it to a global variable and use it in the class.
Potential issues:
1. `load_inline` compiles on first import. It might take time. That's fine.
2. The compilation environment might not allow `-O3`? Usually it's fine. But we can include `extra_cuda_cflags=["-O3"]`.
3. Since we use `torch::Tensor` or `at::Tensor` in the `.cu` file, we need to include `<torch/extension.h>`. That's fine.
4. `positions.data_ptr<int>()` returns a pointer to the int data. If positions is int64, this is wrong. But we convert to int32 in Python. That's safe. We can also add a TORCH_CHECK for dtype:
`TORCH_CHECK(positions.scalar_type() == at::kInt, "positions must be int32");`
Actually, in the host wrapper, we don't need to check everything strictly, but it's good.
One more nuance: The `forward` might be called with inputs on CPU during `run_correctness`? Let's think: The evaluation harness likely tests on GPU because it's a CUDA kernel challenge. But to be robust, we can handle CPU fallback without try/except? The prompt says "Do not use try/except fallbacks to the reference". It doesn't forbid if-statements. But static check might detect if we directly call the reference implementation. To be safe, if the inputs are not CUDA, we can fall back to the reference Python logic (loop). But static check might see that as a fallback. The prompt specifically says: "static-analysis pass that detects reward-hacking patterns (try/except fallbacks to the reference, timing-function patches, lazy-tensor tricks, threading injection, etc.)". It might flag any fallback to reference. So better to not fallback. If inputs are CPU, the kernel will fail. But I think evaluation is on GPU. In many of these tasks, they explicitly test on GPU. Let's assume that.
Another thought: Could we avoid clone of cache by using `torch.empty` for output and then copying with a kernel? The current approach clones both caches and then overwrites some positions. Cloning involves two memcpy operations of size B*cache_len*H*D = 2*128*4*32 = 32768 floats each = 128KB each. That's negligible. But we can also just allocate output tensor and copy caches into the first two slices. Not needed.
Wait, the output is `torch.stack((updated_k, updated_v), dim=0)`. That copies the data again into a new tensor of shape [2, B, ...]. So we have:
- clone cache_k: 128KB
- clone cache_v: 128KB
- stack: copies both into 256KB new tensor.
Total allocation 512KB. Tiny.
Could we allocate the stacked tensor directly and write into it? For example:
out = torch.empty(2, B, cache_len, H, D, dtype=..., device=...)
updated_k = out[0]
updated_v = out[1]
Then copy cache_k into updated_k and cache_v into updated_v. The copy can be done with `updated_k.copy_(cache_k)` and `updated_v.copy_(cache_v)`. That's same as clone + stack. Not a big deal.
Alternatively, we can have the kernel also copy the caches? No, because we only update specific positions. We need to preserve existing cache entries. So we need the old data in the output.
Could we have the kernel write directly to the stacked output and also copy old cache? That would require reading old cache inside the kernel. We could launch a separate kernel for clone, or use cudaMemcpy. But clone is already using cudaMemcpy. So it's fine.
Actually, `torch.stack` also copies. Could we avoid that by returning a tuple? The reference returns a single tensor. So we must return a single tensor. We can allocate `out = torch.empty(2, B, cache_len, H, D, device=cache_k.device, dtype=cache_k.dtype)`, then do `out[0].copy_(cache_k)` and `out[1].copy_(cache_v)`, then pass `out[0]` and `out[1]` to the kernel. This avoids an extra clone + stack, and just does two copies + write. Not a big difference.
But perhaps we can do `out = torch.stack([cache_k, cache_v], dim=0).contiguous(); updated_k = out[0]; updated_v = out[1]; module.rope_kv_update(...)`? That's one stack followed by in-place update. But stack will allocate a new tensor and copy both caches into it. Then the kernel overwrites the relevant parts. That is actually one copy + kernel. Wait, stack copies both into a contiguous buffer. So it's equivalent to clone + stack but done in one operation. But `torch.stack` returns a new tensor, and `out[0]` and `out[1]` are views into it. The kernel can write to those views because they share storage.
Let's do:
out = torch.stack((cache_k, cache_v), dim=0)
updated_k = out[0]
updated_v = out[1]
Then call kernel on updated_k and updated_v.
This avoids two clones and one extra stack copy? Wait, `torch.stack` always copies. So it's one copy of both caches into the output. Then the kernel overwrites some positions. That's efficient.
Let's verify: `out` is contiguous? `torch.stack` creates a new dimension and copies. The resulting tensor is contiguous. `updated_k = out[0]` is a view (not contiguous if we check `is_contiguous`? Actually for a contiguous tensor of shape [2, B, ...], slicing the first dimension yields a tensor with stride [stride_1, stride_2, ...] where the first stride is the product of the rest. Wait, `out.stride()` = [B*cache_len*H*D, cache_len*H*D, H*D, D, 1]. Then `out[0]` has offset 0 and size [B, cache_len, H, D], stride [cache_len*H*D, H*D, D, 1]. That's contiguous! Actually yes, slicing the first dimension of a contiguous tensor yields a contiguous tensor (if the sliced dimension is the first). So `updated_k` is contiguous. So passing its data_ptr to the kernel is fine.
Similarly `out[1]` is contiguous.
So we can do:
out = torch.stack((cache_k, cache_v), dim=0)
updated_k = out[0]
updated_v = out[1]
module.rope_kv_update(...)
return out
But `torch.stack` expects tensors of same shape. They are. And it copies them. This is a single allocation and copy.
Alternatively, we can allocate `out = torch.empty(2, B, cache_len, H, D, device=..., dtype=...)` and then `out[0].copy_(cache_k); out[1].copy_(cache_v);`. That's also one allocation, two copies. `torch.stack` does the same but maybe fused? Not sure. Either is fine.
Let's just use `torch.stack`. It's clean.
But we need to make sure `out[0]` and `out[1]` are contiguous. As argued, because the original tensor is contiguous and we slice the first dimension, the sliced tensors are contiguous. Let's confirm:
If tensor shape (2, B, C, H, D) and stride (B*C*H*D, C*H*D, H*D, D, 1), then tensor[0] has shape (B, C, H, D) and stride (C*H*D, H*D, D, 1). Yes, that's contiguous.
So our Python forward:
def forward(self, k_new, v_new, cos, sin, cache_k, cache_v, positions):
out = torch.stack((cache_k, cache_v), dim=0)
k_new_c = k_new.contiguous()
v_new_c = v_new.contiguous()
cos_c = cos.contiguous()
sin_c = sin.contiguous()
pos_c = positions.to(torch.int32).contiguous()
module.rope_kv_update(k_new_c, v_new_c, cos_c, sin_c, pos_c, out[0], out[1])
return out
This is efficient.
But wait: `torch.stack` might return a tensor that shares memory? Actually `torch.stack` always copies. Let's check PyTorch docs: "Concatenates a sequence of tensors along a new dimension. All tensors need to be of the same size." It creates a new tensor. So yes, it copies.
We should also ensure device: if inputs are on CPU, `.cuda()` call inside forward? No, we shouldn't fallback. But maybe we can do `device = k_new.device` and allocate on that device. `torch.stack` will use the device of the inputs. Since all inputs are on the same device, it's fine.
Wait, but `k_new` and `cache_k` may be on CPU. Then `out` will be on CPU. Then `module.rope_kv_update` expects CUDA. It will throw an error. But as said, evaluation is on GPU.
Now, let's check if the kernel handles the position indexing correctly. `positions` has shape [B, T]. We access `positions[b * T + t]`. That's correct for contiguous row-major.
One subtlety: The cos and sin tables are of shape [cache_len, half_D]. So `cos[pos * half_D + d]` is correct.
What about half_D being zero? D is at least 2 (since 0::2 and 1::2). So half_D >= 1.
Now, we need to compile and test. Let's write the code and use `compile_kernel`. Actually, we have to provide a complete Python file. But for testing, we can write a temporary script and call `compile_kernel`. The tool expects the full kernel file. It will compile and check for errors. Let's draft the code and run compile_kernel.
But before that, we must decide whether to put everything in `cuda_sources` or split. I think putting everything in `cuda_sources` is easiest.
Wait, there is a subtlety: `torch.utils.cpp_extension.load_inline` might require `with_cuda=True` or `cuda_sources` parameter. The function signature:
torch.utils.cpp_extension.load_inline(name, cpp_sources, cuda_sources=None, functions=None, extra_cflags=None, extra_cuda_cflags=None, extra_ldflags=None, extra_include_paths=None, build_directory=None, verbose=False, with_cuda=None, is_python_module=True)
If `cuda_sources` is provided, `with_cuda` is automatically set to True? The docs say: "with_cuda (bool, optional) – determines whether CUDA headers and libraries are added to the build. If set to None (default), this value is automatically determined by checking if cuda_sources is provided."
So we can just provide `cuda_sources`.
Now, the cuda code string. Let's write it carefully.
One more thing: We might want to use `__launch_bounds__(256)` for the kernel? Not necessary.
Also, we might want to handle the case where `B*T*H*half_D` is 0. Not relevant.
Let's write the code.
Potential hazard: `reinterpret_cast<float2*>(updated_k)` when updated_k is a `torch::Tensor` passed to the kernel. In the kernel, we receive `float*`, so it's fine.
Wait, in the kernel signature, should we use `float*` or `const float2*`? We used `float*` for generic pointer and then reinterpret. That's okay.
But we need to make sure `float2` is defined. It's a built-in CUDA vector type. `<cuda_runtime.h>` includes it.
Now, the host wrapper. Since it's in a `.cu` file, we need to include `<torch/extension.h>` which includes pybind11 stuff. But because `load_inline` auto-generates the pybind11 module, we don't need to include it? The pybind11 module is generated in a separate `.cpp` file by `load_inline`. It will call our function. Our `.cu` file just needs to define the function with C++ linkage and the correct signature. It does not need to include pybind11 headers unless we want to use pybind11 types. But we are using `at::Tensor`, which is defined in `<ATen/ATen.h>` or `<torch/extension.h>`. Including `<torch/extension.h>` is safe and provides `at::Tensor`, `TORCH_CHECK`, etc.
However, `torch/extension.h` includes Python.h and pybind11. But since our `.cu` file is compiled by nvcc, and it includes Python.h, there might be warnings or issues. But it's standard practice. Alternatively, we can include `<ATen/ATen.h>` and `<c10/cuda/CUDAGuard.h>` if needed. But `torch/extension.h` is easier.
Let's use `#include <torch/extension.h>`.
One more nuance: The evaluation harness might compile with `-std=c++17`. `load_inline` defaults to C++17. That's fine.
The `extra_cuda_cflags` could include `-O3`. Let's add that.
Now, I will draft the kernel code and use `compile_kernel` to check.
Wait, the `compile_kernel` tool expects the full Python source of the ModelNew kernel file. It does not need the reference code. It will compile the CUDA extension and check for compilation errors. Let's prepare it.
We should define the `ModelNew` class and the `load_inline` call at module level.
Here is a draft:
```python
import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline
CUDA_SOURCE = r"""
#include <torch/extension.h>
#include <cuda_runtime.h>
__global__ void rope_kv_update_kernel(
const float* __restrict__ k_new,
const float* __restrict__ v_new,
const float* __restrict__ cos,
const float* __restrict__ sin,
const int* __restrict__ positions,
float* __restrict__ updated_k,
float* __restrict__ updated_v,
int B,
int T,
int cache_len,
int H,
int D)
{
int half_D = D >> 1;
int total_pairs = B * T * H * half_D;
int idx = blockIdx.x * blockDim.x + threadIdx.x;
if (idx >= total_pairs) return;
int d = idx % half_D;
int tmp = idx / half_D;
int h = tmp % H;
tmp /= H;
int t = tmp % T;
int b = tmp / T;
int pos = positions[b * T + t];
int pair_offset = ((b * T + t) * H + h) * half_D + d;
int cache_pair_offset = ((b * cache_len + pos) * H + h) * half_D + d;
const float2* k_new_ptr = reinterpret_cast<const float2*>(k_new);
const float2* v_new_ptr = reinterpret_cast<const float2*>(v_new);
float2* updated_k_ptr = reinterpret_cast<float2*>(updated_k);
float2* updated_v_ptr = reinterpret_cast<float2*>(updated_v);
float2 x = k_new_ptr[pair_offset];
float c = cos[pos * half_D + d];
float s = sin[pos * half_D + d];
float2 rot;
rot.x = x.x * c - x.y * s;
rot.y = x.x * s + x.y * c;
updated_k_ptr[cache_pair_offset] = rot;
updated_v_ptr[cache_pair_offset] = v_new_ptr[pair_offset];
}
void rope_kv_update(
const at::Tensor& k_new,
const at::Tensor& v_new,
const at::Tensor& cos,
const at::Tensor& sin,
const at::Tensor& positions,
at::Tensor updated_k,
at::Tensor updated_v)
{
TORCH_CHECK(k_new.is_cuda(), "k_new must be a CUDA tensor");
TORCH_CHECK(v_new.is_cuda(), "v_new must be a CUDA tensor");
TORCH_CHECK(cos.is_cuda(), "cos must be a CUDA tensor");
TORCH_CHECK(sin.is_cuda(), "sin must be a CUDA tensor");
TORCH_CHECK(positions.is_cuda(), "positions must be a CUDA tensor");
TORCH_CHECK(updated_k.is_cuda(), "updated_k must be a CUDA tensor");
TORCH_CHECK(updated_v.is_cuda(), "updated_v must be a CUDA tensor");
int B = k_new.size(0);
int T = k_new.size(1);
int H = k_new.size(2);
int D = k_new.size(3);
int cache_len = updated_k.size(1);
TORCH_CHECK(D % 2 == 0, "head_dim must be even");
TORCH_CHECK(positions.scalar_type() == at::kInt, "positions must be int32");
int half_D = D / 2;
int total_pairs = B * T * H * half_D;
const int threads = 256;
const int blocks = (total_pairs + threads - 1) / threads;
rope_kv_update_kernel<<<blocks, threads>>>(
k_new.data_ptr<float>(),
v_new.data_ptr<float>(),
cos.data_ptr<float>(),
sin.data_ptr<float>(),
positions.data_ptr<int>(),
updated_k.data_ptr<float>(),
updated_v.data_ptr<float>(),
B, T, cache_len, H, D);
}
"""
module = load_inline(
name="rope_kv_cache",
cpp_sources="",
cuda_sources=CUDA_SOURCE,
functions=["rope_kv_update"],
extra_cuda_cflags=["-O3"],
verbose=False,
)
class ModelNew(nn.Module):
def __init__(self):
super().__init__()
def forward(
self,
k_new: torch.Tensor,
v_new: torch.Tensor,
cos: torch.Tensor,
sin: torch.Tensor,
cache_k: torch.Tensor,
cache_v: torch.Tensor,
positions: torch.Tensor,
) -> torch.Tensor:
out = torch.stack((cache_k, cache_v), dim=0)
k_new_c = k_new.contiguous()
v_new_c = v_new.contiguous()
cos_c = cos.contiguous()
sin_c = sin.contiguous()
pos_c = positions.to(torch.int32).contiguous()
module.rope_kv_update(k_new_c, v_new_c, cos_c, sin_c, pos_c, out[0], out[1])
return out
```
Wait, will `out[0]` and `out[1]` be contiguous? As discussed, for a contiguous `out`, slicing first dimension yields contiguous subtensors. Let's double-check with a small example mentally.
If `out` is shape (2, B, C, H, D) and stride (B*C*H*D, C*H*D, H*D, D, 1). `out[0]` has offset 0, shape (B, C, H, D), stride (C*H*D, H*D, D, 1). For a tensor of shape (B, C, H, D), the strides are (C*H*D, H*D, D, 1). This is the standard row-major contiguous stride. So yes, it's contiguous.
Even if it weren't, we could call `.contiguous()` on the slices, but that defeats the purpose. But we don't need to.
However, there is a subtlety: `torch.stack` might not produce a tensor with exactly those strides if the inputs are not contiguous? Actually `torch.stack` always produces a new contiguous tensor. It copies the data. So the output is contiguous. Good.
But what if `cache_k` is on CPU? `torch.stack` still produces a CPU tensor. The kernel fails. But as I said, assume GPU.
Now, one thing: The reference uses `cache_k.clone()` and `cache_v.clone()`. Our approach uses `torch.stack`. Is this numerically equivalent? Yes, because we preserve all cache values and update the relevant positions. The output tensor shape and values are the same.
One more nuance: The reference returns a new tensor that is the stack. Our `out` is the stack. Good.
Now, let's think about the `positions` dtype. In `get_inputs`, it's `torch.int32`. The reference does `int(positions[b, t].item())`. Our kernel reads as `int*`. If we do `.to(torch.int32)`, it ensures that. Good.
One more nuance: The cos and sin tables are generated as float32. Good.
Now, compile check. Let's call the tool.
But before that, maybe there is a better way to handle the caching: The reference does `updated_k = cache_k.clone()` then writes. Our `torch.stack` copies as well. But `torch.stack` will allocate memory and copy both. The kernel overwrites some positions. That is correct.
Now, is there any potential issue with `load_inline` recompiling every time? No, it caches.
I think this is ready. Let's run compile_kernel.