reasoning
Let me analyze this task. I need to create a custom CUDA kernel that replaces the PyTorch reference implementation for a low-rank expansion of compressed KV activations.
The reference implementation does:
1. Reshape hidden from (batch_size, seq_len, model_dim) to (batch_size * seq_len, model_dim)
2. Compute latent = hidden @ w_down -> shape (batch_size * seq_len, rank)
3. Compute k = latent @ w_up_k -> shape (batch_size * seq_len, kv_heads * head_dim)
4. Compute v = latent @ w_up_v -> shape (batch_size * seq_len, kv_heads * head_dim)
5. Reshape k to (batch_size, seq_len, kv_heads, head_dim)
6. Reshape v to (batch_size, seq_len, kv_heads, head_dim)
7. Stack k and v along dim=2 -> shape (batch_size, seq_len, 2, kv_heads, head_dim)
Given dimensions:
- batch_size = 4
- seq_len = 96
- model_dim = 256
- rank = 48
- kv_heads = 4
- head_dim = 32
So the matrix multiplications are:
1. hidden: (384, 256) @ w_down: (256, 48) -> latent: (384, 48)
2. latent: (384, 48) @ w_up_k: (48, 128) -> k: (384, 128)
3. latent: (384, 48) @ w_up_v: (48, 128) -> v: (384, 128)
These are relatively small matrix multiplications. The total FLOPs are:
- First matmul: 384 * 256 * 48 = 4.7M FLOPs
- Second matmul (k): 384 * 48 * 128 = 2.4M FLOPs
- Third matmul (v): 384 * 48 * 128 = 2.4M FLOPs
- Total: ~9.5M FLOPs
For such small sizes, cuBLAS is already highly optimized. But I can potentially gain performance by fusing all operations into a single kernel, which eliminates intermediate memory transfers and the separate reshape/stack steps. I can also optimize memory access patterns and leverage shared memory for matrix tiles.
The output tensor has shape (4, 96, 2, 4, 32) = 98,304 elements, with hidden at 98,304 elements, w_down at 12,288, and each weight matrix at 6,144 elements. The arithmetic intensity suggests this could be either memory or compute bound depending on implementation. I should explore fusing the three matrix multiplications into a single kernel launch to reduce kernel launch overhead and eliminate intermediate results storage. I'll start with a basic fused implementation and use profiling to identify optimization opportunities.
My approach will be a tiled matrix multiplication kernel where I load hidden and w_down tiles into shared memory, compute partial results for latent, then immediately multiply with w_up_k and w_up_v tiles without writing latent to global memory. This requires carefully managing shared memory usage for the weight matrices.
Given the dimensions—M=384, K1=256, N1=48, N2=128—I need to calculate shared memory requirements. The hidden matrix needs 384×256 floats, w_down needs 256×48 floats, and each of w_up_k and w_up_v needs 48×128 floats. Since the full matrices won't fit in shared memory at once, I'll tile the computation across multiple kernel launches or iterations. ), I can implement a fused kernel that processes multiple outputs per thread block.
A practical approach would be to compute the latent vector in registers or shared memory for each tile of the output, then immediately apply the up-projection weights without round-tripping to global memory. Each thread block could handle a tile of the M dimension (batch*seq) and compute both k and v outputs simultaneously by reusing the same latent values.
Let me try a straightforward CUDA kernel using cuBLAS for the matrix multiplications first, then consider whether fusing the operations would provide meaningful speedup given the problem sizes. Alternatively, I could write a custom kernel where each thread processes one output element of both k and v, accumulating the necessary intermediate results in registers to avoid redundant computation. in the inner loop)
3. For each k element, compute: sum over r of latent[r] * w_up_k[r, j]
4. For each v element, compute: sum over r of latent[r] * w_up_v[r, j]
This requires iterating over the K dimensions (model_dim and rank) for each output element. A tiled approach would be more efficient.
I'll structure this as:
- Thread blocks process tiles of the (batch*seq, kv_heads*head_dim) output space
- Shared memory loads tiles of hidden, w_down, w_up_k, and w_up_v
- Threads accumulate partial sums and write final results
Let me code this up:
```python
import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline
cuda_source = """
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
// Tiled fused GEMM kernel
// Computes: k = (hidden @ w_down) @ w_up_k, v = (hidden @ w_down) @ w_up_v
// Then reshapes and stacks
#define TILE_SIZE 16
__global__ void fused_kv_kernel(
const float* __restrict__ hidden,
const float* __restrict__ w_down,
const float* __restrict__ w_up_k,
const float* __restrict__ w_up_v,
float* __restrict__ output,
int M, int K1, int N1, int N2) {
// M = batch_size * seq_len
// K1 = model_dim
// N1 = rank
// N2 = kv_heads * head_dim
int row = blockIdx.y * TILE_SIZE + threadIdx.y; // M dimension
int col = blockIdx.x * TILE_SIZE + threadIdx.x; // N2 dimension
__shared__ float hidden_tile[TILE_SIZE][TILE_SIZE];
__shared__ float w_down_tile[TILE_SIZE][TILE_SIZE];
__shared__ float w_up_k_tile[TILE_SIZE][TILE_SIZE];
__shared__ float w_up_v_tile[TILE_SIZE][TILE_SIZE];
float latent_reg[TILE_SIZE] = {0};
float k_val = 0.0f;
float v_val = 0.0f;
// First GEMM: latent = hidden @ w_down
// Each thread computes TILE_SIZE elements of latent (along N1 dimension)
for (int t = 0; t < (K1 + TILE_SIZE - 1) / TILE_SIZE; ++t) {
// Load tiles
int k_idx = t * TILE_SIZE + threadIdx.x;
if (row < M && k_idx < K1) {
hidden_tile[threadIdx.y][threadIdx.x] = hidden[row * K1 + k_idx];
} else {
hidden_tile[threadIdx.y][threadIdx.x] = 0.0f;
}
int n1_idx = t * TILE_SIZE + threadIdx.y;
if (n1_idx < N1 && k_idx < K1) {
w_down_tile[threadIdx.x][threadIdx.y] = w_down[k_idx * N1 + n1_idx];
} else {
w_down_tile[threadIdx.x][threadIdx.y] = 0.0f;
}
__syncthreads();
// Accumulate partial sums for latent
for (int k = 0; k < TILE_SIZE; ++k) {
latent_reg[threadIdx.x] += hidden_tile[threadIdx.y][k] * w_down_tile[k][threadIdx.x];
}
__syncthreads();
}
// Now latent_reg contains partial results for latent[row, t*TILE_SIZE + threadIdx.x]
// We need to complete the reduction across thread blocks
// Actually this is getting complicated. Let me simplify.
}
```
Actually, this tiled approach is getting complex. Let me try a simpler approach first - use cuBLAS for the matmuls and just do the reshaping/stacking in a custom kernel.
But actually, for better performance, let me try a different approach: compute all three matmuls using cuBLAS (which is highly optimized) and then use a custom kernel for the reshape and stack operations.
Wait, but the reshape and stack are just view operations in PyTorch - they don't actually move data, they just change the tensor metadata. The stack operation does require creating a new tensor and copying data.
Let me look at the output format more carefully:
- k has shape (batch_size, seq_len, kv_heads, head_dim) = (4, 96, 4, 32)
- v has shape (batch_size, seq_len, kv_heads, head_dim) = (4, 96, 4, 32)
- output = torch.stack((k, v), dim=2) has shape (4, 96, 2, 4, 32)
The stack operation interleaves k and v along a new dimension at position 2. So output[b, s, 0, h, d] = k[b, s, h, d] and output[b, s, 1, h, d] = v[b, s, h, d].
For maximum performance, I should write a custom kernel that:
1. Computes latent = hidden @ w_down
2. Computes k = latent @ w_up_k
3. Computes v = latent @ w_up_v
4. Writes k and v to the output in the stacked format
Actually, let me try using cublasGemmStridedBatchedEx or just regular cublasSgemm for the matmuls, and then write the output directly in the stacked format.
Actually, the simplest approach would be:
1. Use cublasSgemm to compute latent
2. Use cublasSgemm to compute k and v
3. Write a small kernel to reshape and stack
But this requires 3 kernel launches and storing intermediate latent.
A more efficient approach:
1. Use cublasSgemm to compute latent
2. Use a fused kernel that computes both k = latent @ w_up_k and v = latent @ w_up_v and writes them in stacked format
Or even better:
1. Use a fully fused custom kernel
Let me try the fully fused approach with a simpler implementation:
```python
import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline
cuda_source = """
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <cublas_v2.h>
#define BLOCK_SIZE 256
__global__ void fused_kv_output_kernel(
const float* __restrict__ k_flat,
const float* __restrict__ v_flat,
float* __restrict__ output,
int batch_size, int seq_len, int kv_heads, int head_dim) {
int idx = blockIdx.x * blockDim.x + threadIdx.x;
int total = batch_size * seq_len * kv_heads * head_dim;
if (idx < total) {
// k_flat and v_flat are in (batch*seq, heads*head_dim) format
// output is in (batch, seq, 2, heads, head_dim) format
int head_dim_idx = idx % head_dim;
int heads_idx = (idx / head_dim) % kv_heads;
int seq_idx = (idx / (head_dim * kv_heads)) % seq_len;
int batch_idx = idx / (head_dim * kv_heads * seq_len);
// Output indices
// output[batch, seq, 0, heads, head_dim] = k
// output[batch, seq, 1, heads, head_dim] = v
int out_k_idx = ((batch_idx * seq_len + seq_idx) * 2 + 0) * kv_heads * head_dim + heads_idx * head_dim + head_dim_idx;
int out_v_idx = ((batch_idx * seq_len + seq_idx) * 2 + 1) * kv_heads * head_dim + heads_idx * head_dim + head_dim_idx;
output[out_k_idx] = k_flat[idx];
output[out_v_idx] = v_flat[idx];
}
}
torch::Tensor fused_kv_forward(
torch::Tensor hidden,
torch::Tensor w_down,
torch::Tensor w_up_k,
torch::Tensor w_up_v) {
// hidden: (batch_size, seq_len, model_dim)
// w_down: (model_dim, rank)
// w_up_k: (rank, kv_heads * head_dim)
// w_up_v: (rank, kv_heads * head_dim)
int batch_size = hidden.size(0);
int seq_len = hidden.size(1);
int model_dim = hidden.size(2);
int rank = w_down.size(1);
int kv_heads_dim = w_up_k.size(1); // kv_heads * head_dim
int kv_heads = 4;
int head_dim = 32;
int M = batch_size * seq_len;
int K1 = model_dim;
int N1 = rank;
int N2 = kv_heads_dim;
// Reshape hidden to (M, model_dim)
auto hidden_flat = hidden.view({M, model_dim});
// Create output tensors
auto latent = torch::empty({M, rank}, hidden.options());
auto k_flat = torch::empty({M, kv_heads_dim}, hidden.options());
auto v_flat = torch::empty({M, kv_heads_dim}, hidden.options());
// Create cuBLAS handle
cublasHandle_t handle;
cublasCreate(&handle);
float alpha = 1.0f;
float beta = 0.0f;
// latent = hidden_flat @ w_down
cublasSgemm(handle, CUBLAS_OP_N, CUBLAS_OP_N,
N1, M, K1,
&alpha,
w_down.data_ptr<float>(), N1,
hidden_flat.data_ptr<float>(), K1,
&beta,
latent.data_ptr<float>(), N1);
// k_flat = latent @ w_up_k
cublasSgemm(handle, CUBLAS_OP_N, CUBLAS_OP_N,
N2, M, N1,
&alpha,
w_up_k.data_ptr<float>(), N2,
latent.data_ptr<float>(), N1,
&beta,
k_flat.data_ptr<float>(), N2);
// v_flat = latent @ w_up_v
cublasSgemm(handle, CUBLAS_OP_N, CUBLAS_OP_N,
N2, M, N1,
&alpha,
w_up_v.data_ptr<float>(), N2,
latent.data_ptr<float>(), N1,
&beta,
v_flat.data_ptr<float>(), N2);
cublasDestroy(handle);
// Create output tensor
auto output = torch::empty({batch_size, seq_len, 2, kv_heads, head_dim}, hidden.options());
// Reshape and stack
int total = M * kv_heads_dim;
int blocks = (total + BLOCK_SIZE - 1) / BLOCK_SIZE;
fused_kv_output_kernel<<<blocks, BLOCK_SIZE>>>(
k_flat.data_ptr<float>(),
v_flat.data_ptr<float>(),
output.data_ptr<float>(),
batch_size, seq_len, kv_heads, head_dim);
return output;
}
"""
model_new_source = """
import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline
cuda_source = \"\"\"
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <cublas_v2.h>
#define BLOCK_SIZE 256
__global__ void fused_kv_output_kernel(
const float* __restrict__ k_flat,
const float* __restrict__ v_flat,
float* __restrict__ output,
int batch_size, int seq_len, int kv_heads, int head_dim) {
int idx = blockIdx.x * blockDim.x + threadIdx.x;
int total = batch_size * seq_len * kv_heads * head_dim;
if (idx < total) {
int head_dim_idx = idx % head_dim;
int heads_idx = (idx / head_dim) % kv_heads;
int seq_idx = (idx / (head_dim * kv_heads)) % seq_len;
int batch_idx = idx / (head_dim * kv_heads * seq_len);
int out_k_idx = ((batch_idx * seq_len + seq_idx) * 2 + 0) * kv_heads * head_dim + heads_idx * head_dim + head_dim_idx;
int out_v_idx = ((batch_idx * seq_len + seq_idx) * 2 + 1) * kv_heads * head_dim + heads_idx * head_dim + head_dim_idx;
output[out_k_idx] = k_flat[idx];
output[out_v_idx] = v_flat[idx];
}
}
torch::Tensor fused_kv_forward(
torch::Tensor hidden,
torch::Tensor w_down,
torch::Tensor w_up_k,
torch::Tensor w_up_v) {
int batch_size = hidden.size(0);
int seq_len = hidden.size(1);
int model_dim = hidden.size(2);
int rank = w_down.size(1);
int kv_heads_dim = w_up_k.size(1);
int kv_heads = 4;
int head_dim = 32;
int M = batch_size * seq_len;
int K1 = model_dim;
int N1 = rank;
int N2 = kv_heads_dim;
auto hidden_flat = hidden.view({M, model_dim});
auto latent = torch::empty({M, rank}, hidden.options());
auto k_flat = torch::empty({M, kv_heads_dim}, hidden.options());
auto v_flat = torch::empty({M, kv_heads_dim}, hidden.options());
cublasHandle_t handle;
cublasCreate(&handle);
float alpha = 1.0f;
float beta = 0.0f;
cublasSgemm(handle, CUBLAS_OP_N, CUBLAS_OP_N,
N1, M, K1,
&alpha,
w_down.data_ptr<float>(), N1,
hidden_flat.data_ptr<float>(), K1,
&beta,
latent.data_ptr<float>(), N1);
cublasSgemm(handle, CUBLAS_OP_N, CUBLAS_OP_N,
N2, M, N1,
&alpha,
w_up_k.data_ptr<float>(), N2,
latent.data_ptr<float>(), N1,
&beta,
k_flat.data_ptr<float>(), N2);
cublasSgemm(handle, CUBLAS_OP_N, CUBLAS_OP_N,
N2, M, N1,
&alpha,
w_up_v.data_ptr<float>(), N2,
latent.data_ptr<float>(), N1,
&beta,
v_flat.data_ptr<float>(), N2);
cublasDestroy(handle);
auto output = torch::empty({batch_size, seq_len, 2, kv_heads, head_dim}, hidden.options());
int total = M * kv_heads_dim;
int blocks = (total + BLOCK_SIZE - 1) / BLOCK_SIZE;
fused_kv_output_kernel<<<blocks, BLOCK_SIZE>>>(
k_flat.data_ptr<float>(),
v_flat.data_ptr<float>(),
output.data_ptr<float>(),
batch_size, seq_len, kv_heads, head_dim);
return output;
}
\"\"\"
cpp_source = \"\"\"
torch::Tensor fused_kv_forward(torch::Tensor hidden, torch::Tensor w_down, torch::Tensor w_up_k, torch::Tensor w_up_v);
\"\"\"
fused_kv_module = load_inline(
name='fused_kv_module',
cpp_sources=cpp_source,
cuda_sources=cuda_source,
functions=['fused_kv_forward'],
extra_cuda_cflags=['-O2'],
verbose=True
)
class ModelNew(nn.Module):
def __init__(self):
super().__init__()
def forward(self, hidden, w_down, w_up_k, w_up_v):
return fused_kv_module.fused_kv_forward(hidden, w_down, w_up_k, w_up_v)
"""
Actually, I realize I should check the memory layout more carefully. In PyTorch, tensors are row-major by default, but cuBLAS assumes column-major layout. So I need to be careful about the transposition.
In PyTorch:
- hidden is (batch, seq, model_dim), stored row-major as hidden[batch][seq][model_dim]
- w_down is (model_dim, rank), stored row-major as w_down[model_dim][rank]
- The matmul hidden @ w_down gives (batch, seq, rank)
For cuBLAS with column-major layout:
- If I want to compute C = A @ B where A is (M, K) row-major and B is (K, N) row-major
- In column-major, A is (K, M) and B is (N, K), and I want C to be (N, M) column-major which equals (M, N) row-major
- So I need to compute C^T = B^T @ A^T
- cublasSgemm with transa=N, transb=N: C = alpha * A * B + beta * C
- If I pass A as (K, M), B as (N, K), then C will be (N, M)
- But A is (M, K) row-major, which is the same as (K, M) column-major!
- So I should use transa=T, transb=N: C = A^T * B
Actually, this is getting confusing. Let me think more carefully.
PyTorch tensors are contiguous in memory with row-major order by default. For a tensor of shape (M, K), the memory layout is:
- element [i][j] is at position i*K + j
cuBLAS expects column-major order. For a matrix of shape (M, K) in column-major:
- element (i, j) is at position i + j*M (note: row, column indexing)
So a PyTorch tensor of shape (M, K) is equivalent to a cuBLAS matrix of shape (K, M) where we interpret the memory differently.
For matrix multiplication C = A @ B:
- A is (M, K) in PyTorch
- B is (K, N) in PyTorch
- C is (M, N) in PyTorch
In cuBLAS terms:
- A_pytorch (M, K) is stored as A_cublas (K, M) where the leading dimension is M
- B_pytorch (K, N) is stored as B_cublas (N, K) where the leading dimension is K
- C_pytorch (M, N) should be stored as C_cublas (N, M) where the leading dimension is M
So I want to compute:
- C_cublas (N, M) = A_cublas^T (M, K) @ B_cublas^T (K, N)
- In cuBLAS: cublasSgemm with transa=T, transb=T, but that would compute C = A^T @ B^T
Actually, let me think again. In cuBLAS, for cublasSgemm(handle, transa, transb, m, n, k, alpha, A, lda, B, ldb, beta, C, ldc):
- C is (m, n) in column-major
- A is (m, k) if transa=N, or (k, m) if transa=T
- B is (k, n) if transb=N, or (n, k) if transb=T
For PyTorch C = A @ B where A is (M, K) and B is (K, N):
- I want C to be (M, N) in PyTorch, which is (N, M) in cuBLAS column-major
- A in PyTorch (M, K) is A in cuBLAS with lda=M, and the matrix shape is (K, M) if we interpret column-major
- B in PyTorch (K, N) is B in cuBLAS with ldb=K, and the matrix shape is (N, K) if we interpret column-major
So in cuBLAS:
- m = N, n = M (for output C)
- A should contribute M rows, B should contribute N columns
- Hmm, this is getting confusing
Let me use a different approach. I'll use the formula:
- C = A @ B in PyTorch
- If A is (M, K) and B is (K, N), then C is (M, N)
- In cuBLAS column-major: cublasSgemm(handle, CUBLAS_OP_T, CUBLAS_OP_N, N, M, K, alpha, B, K, A, M, beta, C, N)
- This computes: C (N, M column-major) = B^T (N, K) @ A (K, M) = B^T @ A
- But I want A @ B, not B^T @ A
Let me try another way:
- cublasSgemm(handle, CUBLAS_OP_N, CUBLAS_OP_T, N, M, K, alpha, B, N, A, K, beta, C, N)
- A is stored as (K, M) with lda=K, and we don't transpose, so it contributes K rows to each column of C
- B is stored as (N, K) with ldb=N, and we transpose it, so it contributes K columns to each row of C
- C (N, M) = B^T @ A^T? No, this is still wrong.
OK let me just use the standard formula for PyTorch row-major to cuBLAS column-major conversion.
For C = A @ B in PyTorch:
- A is (M, K), B is (K, N), C is (M, N)
- In cuBLAS: cublasSgemm(handle, CUBLAS_OP_N, CUBLAS_OP_T, N, M, K, alpha, B, N, A, M, beta, C, N)
- m = N, n = M, k = K
- op(A) is (K, M) with lda = N... wait, lda should be the leading dimension of A in memory
I think I should just use cublasSgemm with row-major flag or just use torch::matmul inside the C++ code.
Actually, let me simplify and just use LibTorch's built-in matmul which handles all this correctly:
```cpp
auto latent = torch::matmul(hidden_flat, w_down);
auto k_flat = torch::matmul(latent, w_up_k);
auto v_flat = torch::matmul(latent, w_up_v);
```
This is simpler and LibTorch will use cuBLAS internally with correct handling of layouts.
But wait, then I'm basically just calling PyTorch functions from C++, which might not be faster.
Let me try a different approach: use a fully custom CUDA kernel that fuses all operations.
Here's the key insight: the intermediate latent tensor is (M, N1) = (384, 48). This is small enough that we could potentially keep parts of it in shared memory or registers.
Let me try a tiled approach where each thread block computes a tile of the output:
- Each block handles a tile of (M, N2) for both k and v outputs
- The block loads the corresponding tile of hidden, all of w_down, and the relevant columns of w_up_k and w_up_v
- The block computes latent for its M tile, then computes k and v
But this requires w_down (256*48 = 12KB) to fit in shared/registers, plus parts of w_up_k and w_up_v.
Actually, the sizes are:
- w_down: 256 * 48 * 4 = 49KB
- w_up_k: 48 * 128 * 4 = 24KB
- w_up_v: 48 * 128 * 4 = 24KB
These won't fit in shared memory (typically 48KB-164KB per SM). So we can't load the full weight matrices.
Let me try a tiled approach where we iterate over tiles of the K dimensions:
```cpp
__global__ void fused_kernel(
const float* hidden, const float* w_down,
const float* w_up_k, const float* w_up_v,
float* output,
int M, int K1, int N1, int N2,
int batch_size, int seq_len, int kv_heads, int head_dim) {
// Each thread computes one element of k and one element of v
// For a given (m, n) output position, we need:
// latent[m, :] = sum over k1 of hidden[m, k1] * w_down[k1, :]
// k[m, n] = sum over n1 of latent[m, n1] * w_up_k[n1, n]
// v[m, n] = sum over n1 of latent[m, n1] * w_up_v[n1, n]
// Combined: k[m, n] = sum over k1, n1 of hidden[m, k1] * w_down[k1, n1] * w_up_k[n1, n]
int m = blockIdx.y * blockDim.y + threadIdx.y; // 0 to M-1
int n = blockIdx.x * blockDim.x + threadIdx.x; // 0 to N2-1
if (m >= M || n >= N2) return;
float k_val = 0.0f;
float v_val = 0.0f;
for (int k1 = 0; k1 < K1; k1++) {
float hidden_val = hidden[m * K1 + k1];
for (int n1 = 0; n1 < N1; n1++) {
float latent_val = hidden_val * w_down[k1 * N1 + n1];
k_val += latent_val * w_up_k[n1 * N2 + n];
v_val += latent_val * w_up_v[n1 * N2 + n];
}
}
// Write output in stacked format
// output shape: (batch_size, seq_len, 2, kv_heads, head_dim)
// We have k_val and v_val for position (m, n) in the flat k/v tensors
int batch_idx = m / seq_len;
int seq_idx = m % seq_len;
int kv_head_idx = n / head_dim;
int head_dim_idx = n % head_dim;
// k output
int k_out_idx = ((batch_idx * seq_len + seq_idx) * 2 + 0) * kv_heads * head_dim + kv_head_idx * head_dim + head_dim_idx;
output[k_out_idx] = k_val;
// v output
int v_out_idx = ((batch_idx * seq_len + seq_idx) * 2 + 1) * kv_heads * head_dim + kv_head_idx * head_dim + head_dim_idx;
output[v_out_idx] = v_val;
}
```
This is a naive kernel with O(K1 * N1) work per thread, which is 256 * 48 = 12288 operations per output element. That's a lot of work, but it might still be faster due to no intermediate memory access.
However, this approach has poor memory access patterns - it's reading w_down and w_up_k/w_up_v in a non-coalesced manner.
Let me try a better tiled approach:
Actually, let me step back and think about this more carefully. The matmul sizes are:
1. hidden (384, 256) @ w_down (256, 48) -> latent (384, 48)
2. latent (384, 48) @ w_up_k (48, 128) -> k (384, 128)
3. latent (384, 48) @ w_up_v (48, 128) -> v (384, 128)
The key observation is that matmuls 2 and 3 share the same input (latent). So if I can avoid writing latent to memory, I can save memory bandwidth.
One approach is to use a two-phase kernel:
- Phase 1: Each thread block computes a partial latent for a tile of the M dimension
- Phase 2: Each thread block uses its partial latent to compute k and v for the same tile
But this still requires coordinating between thread blocks for the full latent.
Another approach is to use persistent threads:
- Each thread block is assigned to a specific (m, n) output position
- The block iterates over tiles of K1 and N1, accumulating partial sums
This is similar to my naive kernel above but with better tiling.
Actually, let me try a simpler approach first: use cuBLAS for the three matmuls, but with a custom kernel to reshape and stack the output. This should be reasonably fast and I can profile it to see if it's worth optimizing further.
Let me also consider using torch's built-in functions. The reshape and stack operations in PyTorch might be creating unnecessary copies. Let me check what the reference implementation actually does:
```python
latent = hidden.reshape(batch_size * seq_len, model_dim) @ w_down
k = latent @ w_up_k
v = latent @ w_up_v
k = k.view(batch_size, seq_len, kv_heads, head_dim)
v = v.view(batch_size, seq_len, kv_heads, head_dim)
return torch.stack((k, v), dim=2)
```
- `hidden.reshape()` creates a view (no copy) if possible
- `latent = hidden.reshape(...) @ w_down` creates a new tensor for latent
- `k = latent @ w_up_k` creates a new tensor for k
- `v = latent @ w_up_v` creates a new tensor for v
- `k.view()` creates a view (no copy)
- `v.view()` creates a view (no copy)
- `torch.stack((k, v), dim=2)` creates a new tensor by copying data
So the operations that allocate memory are: latent, k, v, and the final stacked output.
If I fuse the operations, I can avoid allocating latent (and potentially k and v if I write directly to the stacked output).
Let me try a kernel that:
1. Reads hidden and weight matrices
2. Computes latent for a tile of the M dimension
3. Computes k and v for the same tile
4. Writes directly to the stacked output
Here's a tiled implementation:
```cpp
#define TILE_M 16
#define TILE_N 16
#define TILE_K 16
__global__ void fused_kv_tiled_kernel(
const float* __restrict__ hidden,
const float* __restrict__ w_down,
const float* __restrict__ w_up_k,
const float* __restrict__ w_up_v,
float* __restrict__ output,
int M, int K1, int N1, int N2,
int batch_size, int seq_len, int kv_heads, int head_dim) {
__shared__ float hidden_tile[TILE_M][TILE_K];
__shared__ float w_down_tile[TILE_K][TILE_N]; // for N1 dimension
__shared__ float w_up_k_tile[TILE_N][TILE_K]; // for N1 -> N2
__shared__ float w_up_v_tile[TILE_N][TILE_K];
// Each thread block computes TILE_M x TILE_N output elements for both k and v
int m_start = blockIdx.y * TILE_M;
int n_start = blockIdx.x * TILE_N;
int m_local = threadIdx.y; // 0 to TILE_M-1
int n_local = threadIdx.x; // 0 to TILE_N-1
int m = m_start + m_local;
int n = n_start + n_local;
// Accumulators for latent, k, and v
float latent[TILE_N] = {0}; // latent[m, n1_start:n1_start+TILE_N]
float k_val = 0.0f;
float v_val = 0.0f;
// Phase 1: Compute latent = hidden @ w_down
// Iterate over tiles of K1 (model_dim)
for (int k1_tile = 0; k1_tile < (K1 + TILE_K - 1) / TILE_K; k1_tile++) {
int k1_start = k1_tile * TILE_K;
// Load hidden tile
if (m < M && k1_start + m_local < K1) {
hidden_tile[m_local][n_local] = hidden[m * K1 + k1_start + n_local];
} else {
hidden_tile[m_local][n_local] = 0.0f;
}
// Load w_down tile
// w_down is (K1, N1), we're loading w_down[k1_start:k1_start+TILE_K, :]
// But we want to compute latent[m, :] which requires all of w_down[:, :]
// Actually, let me reconsider the tiling...
__syncthreads();
// ... this is getting complicated
}
}
```
OK, this tiled approach is getting quite complex. Let me simplify and try a different strategy.
Let me use a simpler approach where I use cuBLAS for the matmuls, but optimize the memory allocation and data movement:
```cpp
torch::Tensor fused_kv_forward(
torch::Tensor hidden,
torch::Tensor w_down,
torch::Tensor w_up_k,
torch::Tensor w_up_v) {
// Use contiguous inputs
hidden = hidden.contiguous();
w_down = w_down.contiguous();
w_up_k = w_up_k.contiguous();
w_up_v = w_up_v.contiguous();
// ... use cuBLAS for matmuls
}
```
Actually, I just realized that I should try the simplest possible approach first and then iterate. Let me write a basic implementation that uses cuBLAS correctly and handles the memory layout properly.
For cuBLAS with row-major matrices, the key insight is:
- A row-major matrix of shape (M, K) is the same as a column-major matrix of shape (K, M)
- For C = A @ B where A is (M, K) row-major and B is (K, N) row-major:
- In column-major terms: C_col (N, M) = A_col^T (M, K) @ B_col (K, N)?
- Actually, let me just use cublasSgemm with the correct parameters
For row-major A (M, K) and B (K, N), to compute C (M, N) row-major:
- In cuBLAS column-major, this is computing C^T = B^T @ A^T
- Use: cublasSgemm(handle, CUBLAS_OP_N, CUBLAS_OP_T, N, M, K, &alpha, B, K, A, K, &beta, C, N)
- op(B) is B (K, N) in col-major, which is (N, K) in row-major? No...
OK I think I'm overcomplicating this. Let me use a different approach:
For row-major matrices, if I want to compute C = A @ B:
- I can compute C^T = B^T @ A^T using cuBLAS in the usual column-major way
- Then C is the transpose of the result
So:
- A (M, K) row-major becomes A^T (K, M) column-major with lda = M
- B (K, N) row-major becomes B^T (N, K) column-major with ldb = K
- C (M, N) row-major becomes C^T (N, M) column-major with ldc = M
cublasSgemm(handle, CUBLAS_OP_N, CUBLAS_OP_N, N, M, K, &alpha, B, ldb, A, lda, &beta, C, ldc)
- This computes C = A @ B in column-major
- If A is (K, M) with lda=K, and B is (N, K) with ldb=N, then C is (N, M) with ldc=N
- But we have A^T (K, M) and B^T (N, K), so this is computing C = A^T @ B^T = (B @ A)^T
Hmm, I think I should use:
cublasSgemm(handle, CUBLAS_OP_T, CUBLAS_OP_T, N, M, K, &alpha, B, ldb, A, lda, &beta, C, ldc)
This computes C = B^T @ A^T = (A @ B)^T, which is what we want for row-major!
With:
- A (M, K) row-major, so A^T is (K, M) col-major with lda = M
- B (K, N) row-major, so B^T is (N, K) col-major with ldb = K
- C (M, N) row-major, so C^T is (N, M) col-major with ldc = M
So the call is:
cublasSgemm(handle, CUBLAS_OP_T, CUBLAS_OP_T, N, M, K, &alpha, B, K, A, M, &beta, C, M)
Wait, let me double-check with the cublasSgemm documentation:
- C (m, n) = alpha * op(A) * op(B) + beta * C
- if transa = CUBLAS_OP_T, op(A) = A^T
- if transb = CUBLAS_OP_T, op(B) = B^T
- A is (k, m) if transa=T, with lda >= k
- B is (n, k) if transb=T, with ldb >= n
- C is (m, n), with ldc >= m
So with transa=T, transb=T, m=N, n=M, k=K:
- op(A) = A^T where A is (n, k) = (M, K), so A^T is (K, M)
- op(B) = B^T where B is (k, m) = (K, N), so B^T is (N, K)
- C = alpha * A^T @ B^T + beta * C
- C is (m, n) = (N, M) in column-major, which is (M, N) in row-major
But wait, for our row-major matrices:
- A (the hidden matrix) is (M, K1) = (384, 256) row-major, stored with stride K1=256
- B (the w_down matrix) is (K1, N1) = (256, 48) row-major, stored with stride N1=48
In column-major interpretation:
- A_col would have shape (K1, M) = (256, 384) with leading dimension 384 (if we reinterpret)
- B_col would have shape (N1, K1) = (48, 256) with leading dimension 256
Actually, I think the right way to think about this is:
- A row-major tensor of shape (rows, cols) has elements laid out as: [row0col0, row0col1, ..., row1col0, row1col1, ...]
- This is the same memory layout as a column-major tensor of shape (cols, rows) with leading dimension = rows
So for A (M, K1) row-major:
- Same as column-major A_col of shape (K1, M) with lda = M (the number of rows in original)
- Wait, lda should be the leading dimension, which is the stride between columns, which is M (the original row stride in row-major, which becomes the column stride in col-major)
Actually, I think I have it backwards. Let me reconsider.
For a row-major array of shape (R, C):
- Element [i, j] is at position i * C + j
- The stride between consecutive elements in a row is 1
- The stride between consecutive elements in a column is C
For a column-major array of shape (R, C):
- Element [i, j] is at position i + j * R
- The stride between consecutive elements in a row is R
- The stride between consecutive elements in a column is 1
- The leading dimension (lda) is R (the stride in the row direction)
So if I have a row-major array of shape (R, C):
- Memory position of [i, j] is i * C + j
- If I reinterpret this as column-major with lda = C:
- Memory position of column-major [i', j'] would be i' + j' * C
- To match i * C + j, I need i' + j' * C = i * C + j
- If j' = i (the row index of row-major becomes the column index of col-major)
- And i' = j (the column index of row-major becomes the row index of col-major)
- Then the shape in col-major is (C, R)
So a row-major (R, C) is equivalent to a column-major (C, R) with lda = C.
Now for matmul C = A @ B in row-major:
- A is (M, K) row-major = (K, M) col-major
- B is (K, N) row-major = (N, K) col-major
- C is (M, N) row-major = (N, M) col-major
In col-major, I want to compute C_col = A_col @ B_col, where:
- A_col is (K, M)
- B_col is (N, K)
- C_col should be (N, M) = B_col @ A_col? No, that doesn't work dimension-wise.
Wait, the matmul dimensions don't match in col-major interpretation:
- A_col is (K, M), B_col is (N, K)
- For A_col @ B_col, we need A_col to have K columns and B_col to have K rows
- But A_col has M columns and B_col has N rows
So the reinterpretation doesn't work directly. I need to actually transpose one of the matrices.
In col-major terms:
- If A_col is (K, M) and I want to use it as if it has K rows and M columns
- And B_col is (N, K) and I want to use it as if it has K rows and N columns
- Then A_col^T is (M, K) and B_col is (N, K)
- For A_col^T @ B_col^T, dimensions are (M, K) @ (K, N) = (M, N)
- But I need B_col^T, which is (K, N) from (N, K)
Actually let me think about what I want:
- I want C_row (M, N) = A_row (M, K) @ B_row (K, N)
- In col-major terms: C_col (N, M) = some operation on A_col (K, M) and B_col (N, K)
Using the identity: (A @ B)^T = B^T @ A^T
- C_row^T = B_row^T @ A_row^T
- C_col = B_row^T @ A_row^T
- In col-major: C_col (N, M) = B_col^T @ A_col^T where B_col is (N, K) and A_col is (K, M)
- B_col^T (K, N) @ A_col^T (M, K) ... dimensions don't match
Hmm, I think I'm making an error somewhere. Let me try a different approach.
The key insight is that cublasSgemm computes C = op(A) @ op(B) where A, B, C are all in column-major order.
For row-major matrices, I can use the identity:
- C_row = A_row @ B_row
- C_row^T = B_row^T @ A_row^T
In cuBLAS:
- cublasSgemm computes C = op(A) @ op(B)
- If I set op(A) = B_row^T and op(B) = A_row^T
- Then C = B_row^T @ A_row^T = C_row^T
But wait, in cuBLAS the input matrices A and B are in column-major order. So I need to pass B_row and A_row as if they were column-major.
For B_row (K, N):
- B_row[i, j] is at position i * N + j (row-major)
- If I pass this as a column-major matrix B_col, it would be interpreted as:
- B_col[i', j'] is at position i' + j' * lda
- With lda = N (the stride in the original row-major)
- The shape would be (N, K) in column-major
So B_row (K, N) becomes B_col (N, K) with lda = N.
Similarly, A_row (M, K) becomes A_col (K, M) with lda = M.
Now I want to compute C = B_row^T @ A_row^T using cuBLAS:
- B_row^T is (N, K) row-major, which is (K, N) col-major
- A_row^T is (K, M) row-major, which is (M, K) col-major
But I have B_col (N, K) and A_col (K, M). To get B_row^T and A_row^T:
- B_row^T (N, K) row-major = B_row (K, N) transposed
- In terms of B_col (N, K), B_col^T would be (K, N) col-major
- But B_col^T col-major (K, N) = B_col (N, K) transposed in memory... this is getting complicated
Let me try yet another approach. I'll use the identity:
- For row-major A (M, K) and B (K, N), C = A @ B can be computed as:
- C = cublasSgemm(..., transa=CUBLAS_OP_N, transb=CUBLAS_OP_T, ...)
- Where we interpret A as (M, K) col-major (but it's actually stored as K x M in memory)
- And B as (K, N) col-major (stored as N x K in memory)
OK I think I need to just try it and see. Let me write the code and test it.
Actually, I just realized there's a simpler approach: I can use LibTorch's matmul function which handles all the layout conversion automatically:
```cpp
auto latent = hidden.view({M, model_dim}).matmul(w_down);
auto k_flat = latent.matmul(w_up_k);
auto v_flat = latent.matmul(w_up_v);
```
This should use cuBLAS internally and handle the layout correctly. The question is whether this is faster than pure PyTorch.
Actually, I think the main optimization opportunity here is:
1. Avoid intermediate memory allocation (latent tensor)
2. Fuse the stack operation
Let me try a custom kernel that fuses the three matmuls and the stack.
Actually, given the small sizes (384x256, 256x48, 384x48, 48x128), cuBLLAS might already be very efficient. Let me first try a basic implementation and see what the performance is.
Let me write the code now:
```python
import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline
cuda_source = """
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <cublas_v2.h>
#define BLOCK_SIZE 256
__global__ void stack_kv_kernel(
const float* __restrict__ k_flat,
const float* __restrict__ v_flat,
float* __restrict__ output,
int batch_size, int seq_len, int kv_heads, int head_dim) {
int idx = blockIdx.x * blockDim.x + threadIdx.x;
int total = batch_size * seq_len * kv_heads * head_dim;
if (idx < total) {
int head_dim_idx = idx % head_dim;
int heads_idx = (idx / head_dim) % kv_heads;
int seq_idx = (idx / (head_dim * kv_heads)) % seq_len;
int batch_idx = idx / (head_dim * kv_heads * seq_len);
// output shape: (batch_size, seq_len, 2, kv_heads, head_dim)
// k is at [b, s, 0, h, d]
// v is at [b, s, 1, h, d]
int flat_idx = batch_idx * seq_len * kv_heads * head_dim + seq_idx * kv_heads * head_dim + heads_idx * head_dim + head_dim_idx;
int k_out_idx = ((batch_idx * seq_len + seq_idx) * 2) * kv_heads * head_dim + heads_idx * head_dim + head_dim_idx;
int v_out_idx = ((batch_idx * seq_len + seq_idx) * 2 + 1) * kv_heads * head_dim + heads_idx * head_dim + head_dim_idx;
output[k_out_idx] = k_flat[flat_idx];
output[v_out_idx] = v_flat[flat_idx];
}
}
torch::Tensor fused_kv_forward(
torch::Tensor hidden,
torch::Tensor w_down,
torch::Tensor w_up_k,
torch::Tensor w_up_v) {
int batch_size = hidden.size(0);
int seq_len = hidden.size(1);
int model_dim = hidden.size(2);
int rank = w_down.size(1);
int kv_heads_dim = w_up_k.size(1);
int kv_heads = 4;
int head_dim = 32;
int M = batch_size * seq_len;
// Use torch matmul for simplicity (handles layouts correctly)
auto hidden_flat = hidden.view({M, model_dim});
auto latent = torch::matmul(hidden_flat, w_down);
auto k_flat = torch::matmul(latent, w_up_k);
auto v_flat = torch::matmul(latent, w_up_v);
// Create output tensor and use custom kernel to stack
auto output = torch::empty({batch_size, seq_len, 2, kv_heads, head_dim}, hidden.options());
int total = M * kv_heads_dim;
int blocks = (total + BLOCK_SIZE - 1) / BLOCK_SIZE;
stack_kv_kernel<<<blocks, BLOCK_SIZE>>>(
k_flat.data_ptr<float>(),
v_flat.data_ptr<float>(),
output.data_ptr<float>(),
batch_size, seq_len, kv_heads, head_dim);
return output;
}
"""
cpp_source = """
torch::Tensor fused_kv_forward(torch::Tensor hidden, torch::Tensor w_down, torch::Tensor w_up_k, torch::Tensor w_up_v);
"""
module = load_inline(
name='fused_kv',
cpp_sources=cpp_source,
cuda_sources=cuda_source,
functions=['fused_kv_forward'],
extra_cuda_cflags=['-O2'],
verbose=True
)
class ModelNew(nn.Module):
def __init__(self):
super().__init__()
def forward(self, hidden, w_down, w_up_k, w_up_v):
return module.fused_kv_forward(hidden, w_down, w_up_k, w_up_v)
```
Actually, wait. If I'm using torch::matmul, I'm basically calling PyTorch's implementation from C++. This might not be any faster. The overhead of calling C++ from Python and then calling back into PyTorch might even make it slower.
Let me think about this differently. What are the potential optimization opportunities?
1. **Fused kernel**: Compute latent, k, and v in a single kernel without storing intermediates
2. **Better memory layout**: Write output in the final stacked format directly
3. **Reduced memory allocations**: Avoid allocating latent, k, and v separately
Let me try a fully fused kernel. The challenge is that the K dimension for the first matmul (256) and the K dimension for the second/third matmuls (48) are different, so we can't easily share tiles.
One approach: Each thread computes one element of k and one element of v. For each output element:
- k[m, n] = sum_{k1, n1} hidden[m, k1] * w_down[k1, n1] * w_up_k[n1, n]
- v[m, n] = sum_{k1, n1} hidden[m, k1] * w_down[k1, n1] * w_up_v[n1, n]
This requires O(K1 * N1) = O(256 * 48) = 12288 operations per output element, but with good memory coalescing.
For M=384 and N2=128, total threads = 384 * 128 = 49152. Each thread does 12288 multiply-adds, so total operations = 49152 * 12288 ≈ 600M FLOPs. This seems high... let me recalculate.
Actually, the naive implementation would compute latent first, then k and v:
- latent (384, 48) = hidden (384, 256) @ w_down (256, 48): 384 * 256 * 48 * 2 = 9.4M FLOPs
- k (384, 128) = latent (384, 48) @ w_up_k (48, 128): 384 * 48 * 128 * 2 = 4.7M FLOPs
- v (384, 128) = latent (384, 48) @ w_up_v (48, 128): 384 * 48 * 128 * 2 = 4.7M FLOPs
- Total: ~19M FLOPs
But if I compute each output element by iterating over all K1 and N1:
- k[m, n] = sum_{k1, n1} hidden[m, k1] * w_down[k1, n1] * w_up_k[n1, n]
- This is a 2D reduction over (k1, n1)
- Each output element requires K1 * N1 = 256 * 48 = 12288 multiplies and additions
- Total: 384 * 128 * 12288 * 2 = 1.2B FLOPs
That's way more FLOPs than the optimized approach! That's because I'm recomputing latent * w_up for each output element of k and v, instead of computing latent once and reusing it.
So the fused approach that avoids storing latent actually increases the FLOPs significantly. This is a trade-off between memory bandwidth and compute.
Given that the computation is not that large (~19M FLOPs) and the memory access patterns for the weights are good (sequential access), the original PyTorch implementation might already be close to optimal.
Let me check what optimizations are actually possible:
1. **Memory bandwidth**: The total memory accessed is:
- hidden: 384 * 256 * 4 = 384KB
- w_down: 256 * 48 * 4 = 48KB
- w_up_k: 48 * 128 * 4 = 24KB
- w_up_v: 48 * 128 * 4 = 24KB
- latent: 384 * 48 * 4 = 72KB (written, then read twice)
- k: 384 * 128 * 4 = 192KB (written, then read once for stack)
- v: 384 * 128 * 4 = 192KB (written, then read once for stack)
- output: 384 * 128 * 2 * 4 = 384KB (written once)
Total memory bandwidth: ~1.3MB + intermediates
2. **Latent reuse**: If we can avoid writing latent to memory, we save 72KB of write + 144KB of read (latent is read twice for k and v).
3. **Stack fusion**: If we can write k and v directly in stacked format, we save the intermediate k and v tensors.
Let me try a different fused approach:
For each output element (for both k and v):
1. Compute the corresponding latent element
2. Use that latent element for all output positions that need it
Wait, that's also not optimal. Let me think about this differently.
If each thread block handles a tile of the output (M, N2), it would:
1. Load the corresponding tile of hidden (TILE_M, K1)
2. Load all of w_down (K1, N1) - but this might not fit in shared memory
3. Compute partial latent
4. Load relevant tiles of w_up_k and w_up_v
5. Compute k and v
6. Write output
The problem is w_down (256*48*4 = 48KB) might not fit in shared memory, depending on the GPU. Modern GPUs have 48KB-164KB of shared memory per SM, so it might fit.
Let me try this approach:
```cpp
#define TILE_M 32
#define TILE_N 32
#define TILE_K 8
__global__ void fused_kv_kernel(
const float* __restrict__ hidden,
const float* __restrict__ w_down,
const float* __restrict__ w_up_k,
const float* __restrict__ w_up_v,
float* __restrict__ output,
int M, int K1, int N1, int N2,
int batch_size, int seq_len, int kv_heads, int head_dim) {
// Each thread block handles TILE_M rows and TILE_N columns
// For each row, we compute the full latent (all N1 elements)
// Then we use the latent to compute k and v for TILE_N columns
__shared__ float hidden_tile[TILE_M][TILE_K + 1]; // +1 to avoid bank conflicts
__shared__ float w_down_tile[TILE_K][N1]; // This is too big!
// ...
}
```
The problem is N1 (rank) = 48, which is too large for shared memory if we want to store the full w_down tile.
Let me try a different tiling strategy:
- Tile over K1 and N1 for the latent computation
- Each thread computes multiple latent elements
- Then use these latent elements to compute multiple output elements
Actually, let me simplify and just try a basic fused implementation without worrying too much about optimal tiling. I'll use the observation that we can compute latent and then immediately use it without writing to global memory.
Here's a simpler approach:
- Each thread block handles a tile of the M dimension (rows)
- Each thread in the block handles one column of the output for both k and v
- The block computes latent for its rows, then uses it to compute k and v
```cpp
#define BLOCK_SIZE 128
#define TILE_M 8
__global__ void fused_kv_kernel_simple(
const float* __restrict__ hidden,
const float* __restrict__ w_down,
const float* __restrict__ w_up_k,
const float* __restrict__ w_up_v,
float* __restrict__ output,
int M, int K1, int N1, int N2,
int batch_size, int seq_len, int kv_heads, int head_dim) {
// Block handles TILE_M rows and N2 columns
// Each thread handles one column for each row in the tile
__shared__ float latent_shared[TILE_M][N1]; // TILE_M * N1 * 4 = 8 * 48 * 4 = 1.5KB
__shared__ float hidden_shared[TILE_M][K1]; // TILE_M * K1 * 4 = 8 * 256 * 4 = 8KB
int m_tile = blockIdx.x; // which tile of M
int m_local = threadIdx.y; // which row within tile (0 to TILE_M-1)
int n = threadIdx.x; // which column (0 to N2-1)
int m = m_tile * TILE_M + m_local;
// Load hidden tile into shared memory
if (m < M) {
for (int k1 = 0; k1 < K1; k1 += BLOCK_SIZE / TILE_M) {
int k1_idx = k1 + threadIdx.x;
if (k1_idx < K1) {
hidden_shared[m_local][k1_idx] = hidden[m * K1 + k1_idx];
}
}
}
__syncthreads();
// Compute latent for this row
// Each thread computes one element of latent
// latent[m, n1] = sum_{k1} hidden[m, k1] * w_down[k1, n1]
for (int n1_tile = 0; n1_tile < (N1 + BLOCK_SIZE/TILE_M - 1) / (BLOCK_SIZE/TILE_M); n1_tile++) {
int n1 = n1_tile * (BLOCK_SIZE / TILE_M) + threadIdx.x;
if (m < M && n1 < N1) {
float sum = 0.0f;
for (int k1 = 0; k1 < K1; k1++) {
sum += hidden_shared[m_local][k1] * w_down[k1 * N1 + n1];
}
latent_shared[m_local][n1] = sum;
}
}
__syncthreads();
// Compute k and v for this row and column
// k[m, n] = sum_{n1} latent[m, n1] * w_up_k[n1, n]
// v[m, n] = sum_{n1} latent[m, n1] * w_up_v[n1, n]
if (m < M && n < N2) {
float k_val = 0.0f;
float v_val = 0.0f;
for (int n1 = 0; n1 < N1; n1++) {
k_val += latent_shared[m_local][n1] * w_up_k[n1 * N2 + n];
v_val += latent_shared[m_local][n1] * w_up_v[n1 * N2 + n];
}
// Write to output in stacked format
int batch_idx = m / seq_len;
int seq_idx = m % seq_len;
int kv_head_idx = n / head_dim;
int head_dim_idx = n % head_dim;
int k_out_idx = ((batch_idx * seq_len + seq_idx) * 2) * kv_heads * head_dim + kv_head_idx * head_dim + head_dim_idx;
int v_out_idx = ((batch_idx * seq_len + seq_idx) * 2 + 1) * kv_heads * head_dim + kv_head_idx * head_dim + head_dim_idx;
output[k_out_idx] = k_val;
output[v_out_idx] = v_val;
}
}
```
Actually, this approach has issues. The hidden_shared is too large (TILE_M * K1 = 8 * 256 = 2048 floats = 8KB), and we're loading it with a loop that might not be correct.
Let me revise this with proper tiling:
Actually, I think I should just start with a working implementation and then iterate. Let me write a simpler kernel that:
1. Uses cuBLAS for the three matmuls (fast, well-optimized)
2. Uses a custom kernel to stack k and v into the output format
This avoids the complexity of the fully fused kernel while still providing some speedup from avoiding Python overhead.
Let me write and test this implementation first.