reasoning
reasoning (truncated 63694 chars): This is a complex task - I need to implement a full Mamba block in CUDA. Let me break down the components:
1. RMSNorm
2. Linear projection (in_proj) - expands to d_inner * 2
3. Split into x_branch and z branches
4. Causal Conv1d on x_branch
5. SiLU activation on x_branch
6. Selective scan (SSM) - the core Mamba operation
7. Element-wise gate with z
8. Linear projection (out_proj)
9. Residual connection
The selective scan is the most computationally intensive part - it's a recurrent operation that processes the sequence step by step.
Let me think about how to optimize this:
1. **RMSNorm**: Simple element-wise operation with reduction for computing RMS
2. **Linear projections**: These are matrix multiplications, can be optimized
3. **Conv1d**: Depthwise convolution, can be optimized
4. **Selective scan**: This is sequential by nature but can be parallelized across batch and channel dimensions
For the selective scan:
- For each time step t, we update h = dA[:, t] * h + x_db[:, t]
- Then compute y_t = (h * C_t[:, t].unsqueeze(1)).sum(dim=-1)
- This is inherently sequential in the time dimension but parallel in batch and channel dimensions
The key insight is that we can parallelize across the batch dimension and the d_inner dimension, but the sequence dimension must be processed sequentially (though we could potentially use parallel scan algorithms).
Let me design the CUDA kernels:
1. **Combined kernel approach**: I could try to fuse multiple operations together, but the selective scan is complex enough that I'll focus on optimizing that first.
2. **Memory layout**: The input is [B, L, D], and we need to be careful about memory access patterns.
Let me start with a basic implementation and then optimize:
```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 <cmath>
// RMSNorm kernel
__global__ void rmsnorm_kernel(
const float* __restrict__ input,
const float* __restrict__ weight,
float* __restrict__ output,
int B, int L, int D,
float eps
) {
int b = blockIdx.x;
int l = blockIdx.y;
if (b < B && l < L) {
// Compute sum of squares
float sum_sq = 0.0f;
for (int d = 0; d < D; d++) {
float val = input[b * L * D + l * D + d];
sum_sq += val * val;
}
float rms = sqrtf(sum_sq / D + eps);
float inv_rms = 1.0f / rms;
// Normalize and scale
for (int d = 0; d < D; d++) {
int idx = b * L * D + l * D + d;
output[idx] = input[idx] * inv_rms * weight[d];
}
}
}
// ... more kernels
"""
# Define the extension
mamba_extension = load_inline(
name='mamba_extension',
cpp_sources=[],
cuda_sources=[cuda_source],
functions=['rmsnorm_kernel', ...],
verbose=True
)
```
Actually, let me think more carefully about this. The selective scan is the core operation. In Mamba, it's implemented efficiently using a parallel scan algorithm. But implementing that from scratch is complex.
For a simpler approach, I'll parallelize across batch and channel dimensions, processing the sequence sequentially:
```cuda
// Selective scan kernel - one thread block per (batch, channel) pair
__global__ void selective_scan_kernel(
const float* __restrict__ x_db, // [B, L, d_inner, N]
const float* __restrict__ dA, // [B, L, d_inner, N]
const float* __restrict__ C_t, // [B, L, N]
const float* __restrict__ D, // [d_inner]
const float* __restrict__ x_branch, // [B, L, d_inner]
float* __restrict__ y, // [B, L, d_inner]
int B, int L, int d_inner, int N
) {
int b = blockIdx.y;
int c = blockIdx.x; // channel index
if (b < B && c < d_inner) {
float h[16]; // Assuming N <= 16, use shared/local memory
for (int n = 0; n < N; n++) {
h[n] = 0.0f;
}
for (int t = 0; t < L; t++) {
// Update h
for (int n = 0; n < N; n++) {
int idx = ((b * L + t) * d_inner + c) * N + n;
h[n] = dA[idx] * h[n] + x_db[idx];
}
// Compute y_t
float y_t = 0.0f;
for (int n = 0; n < N; n++) {
int c_idx = (b * L + t) * N + n;
y_t += h[n] * C_t[c_idx];
}
// Add D * x_branch
y_t += D[c] * x_branch[(b * L + t) * d_inner + c];
y[(b * L + t) * d_inner + c] = y_t;
}
}
}
```
Wait, I need to be more careful. The C_t tensor has shape [B, L, N] and is shared across channels. Let me re-read the code.
Looking at the reference code:
- `B_t` has shape [B, L, N] (projected from x_branch)
- `C_t` has shape [B, L, N] (projected from x_branch)
- `dA` has shape [B, L, d_inner, N]
- `dB` has shape [B, L, d_inner, N]
- `x_db` has shape [B, L, d_inner, N]
And the selective scan:
```python
for t in range(L):
h = dA[:, t] * h + x_db[:, t]
y_t = (h * C_t[:, t].unsqueeze(1)).sum(dim=-1)
outputs.append(y_t)
```
So `h` has shape [B, d_inner, N], and `C_t[:, t]` has shape [B, N]. After unsqueeze(1), it becomes [B, 1, N], and broadcasting with h [B, d_inner, N] gives [B, d_inner, N]. The sum over dim=-1 gives [B, d_inner].
So each channel has its own hidden state of size N, and they all share the same C_t values.
Let me rewrite the kernel more carefully. I'll use a simpler approach where I parallelize over batch and channel, and use shared memory for the hidden state.
Actually, looking at this more carefully, I think I should:
1. Use PyTorch's built-in operations for some parts (like matrix multiplications for linear layers)
2. Only write custom CUDA for the selective scan and potentially for other fused operations
But the problem asks to replace the forward computation entirely with custom CUDA. Let me think about what the fastest approach would be.
For optimal performance, I would:
1. Fuse RMSNorm with the first linear projection
2. Fuse the conv1d + activation
3. Implement an optimized selective scan
4. Fuse the final operations
However, given the complexity, let me start with a working implementation that uses PyTorch for the linear layers (since they're well-optimized) and custom CUDA for the selective scan and other element-wise operations.
Wait, looking at the problem again, it says to "replace the forward computation" with custom CUDA. I think this means I can still use PyTorch operations but should have custom CUDA for the core computations.
Let me design a cleaner implementation:
1. Use PyTorch's linear layers for in_proj, proj_B, proj_C, proj_delta, out_proj
2. Custom CUDA for:
- RMSNorm
- Fused conv1d + SiLU
- Selective scan
- Final element-wise operations
Actually, for the fastest implementation, I should try to fuse as much as possible. Let me think about memory bandwidth vs compute:
The input is [16, 2048, 256], which is about 32MB for fp32. The d_inner is 512.
For the selective scan:
- Input x_db: [16, 2048, 512, 16] = ~1.3GB
- Hidden state: [16, 512, 16] = ~0.5MB
The sequential scan processes each time step, updating the hidden state and computing output. This is memory-bandwidth bound on the x_db and dA inputs.
Let me design a kernel that:
1. Loads the A matrix once (it's [d_inner, N] and shared across all batch and time)
2. Streams through the sequence, keeping the hidden state in registers
3. Uses efficient memory access patterns
Actually, let me look at what operations need to be done in the forward pass step by step:
1. **Input**: [B, L, D] = [16, 2048, 256]
2. **RMSNorm**: Compute per-token normalization
3. **in_proj**: Linear [B, L, D] -> [B, L, d_inner * 2] = [B, L, 1024]
4. **Split**: x_branch [B, L, d_inner], z [B, L, d_inner]
5. **Conv1d on x_branch**: [B, L, d_inner] -> [B, L, d_inner] (with causal conv)
6. **SiLU
...[ omitted 63694 characters ]...
ECK at the host level should catch this.
Actually, looking at the reference code again:
```python
def get_init_inputs():
return [d_model, d_state, d_conv, expand]
```
The d_state parameter is passed as a constructor argument, so it could theoretically be different from 16. But the default test uses d_state=16, which is <= N_MAX=16. So my implementation should work for the test cases.
Let me finalize and submit the code.
Wait, I also want to consider the data type. The reference implementation uses fp32, and my kernels assume float. I should make sure the inputs are fp32. If the model is used with fp16 or bf16, the kernels would need to be modified. But for now, let me assume fp32.
Also, I should ensure the tensors are on CUDA before calling the CUDA kernels. PyTorch will raise an error if they're not, so I don't need to handle this explicitly.
Let me write the final version:
```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 <cmath>
#define N_MAX 16
__global__ void rmsnorm_kernel(
const float* __restrict__ input,
const float* __restrict__ weight,
float* __restrict__ output,
int B, int L, int D
) {
int idx = blockIdx.x * blockDim.x + threadIdx.x;
int total = B * L;
if (idx < total) {
float sum_sq = 0.0f;
for (int d = 0; d < D; d++) {
float val = input[idx * D + d];
sum_sq += val * val;
}
float rms_inv = rsqrtf(sum_sq / D + 1e-5f);
for (int d = 0; d < D; d++) {
output[idx * D + d] = input[idx * D + d] * rms_inv * weight[d];
}
}
}
__global__ void selective_scan_kernel(
const float* __restrict__ x_branch,
const float* __restrict__ delta,
const float* __restrict__ B_t,
const float* __restrict__ C_t,
const float* __restrict__ A_log,
const float* __restrict__ D_param,
float* __restrict__ output,
int B, int L, int d_inner, int N
) {
int b = blockIdx.y;
int c = blockIdx.x * blockDim.x + threadIdx.x;
if (b >= B || c >= d_inner) return;
float A[N_MAX];
for (int n = 0; n < N; n++) {
A[n] = -expf(A_log[c * N + n]);
}
float h[N_MAX];
for (int n = 0; n < N; n++) {
h[n] = 0.0f;
}
float D_val = D_param[c];
int base_x = b * L * d_inner + c;
int base_BC = b * L * N;
for (int t = 0; t < L; t++) {
float delta_t = delta[base_x + t * d_inner];
float x_t = x_branch[base_x + t * d_inner];
const float* B_t_ptr = B_t + base_BC + t * N;
const float* C_t_ptr = C_t + base_BC + t * N;
float y_t = 0.0f;
for (int n = 0; n < N; n++) {
float dA_t = expf(delta_t * A[n]);
float dB_t = delta_t * B_t_ptr[n];
h[n] = dA_t * h[n] + x_t * dB_t;
y_t += h[n] * C_t_ptr[n];
}
y_t += D_val * x_t;
output[base_x + t * d_inner] = y_t;
}
}
__global__ void gate_kernel(
const float* __restrict__ y,
const float* __restrict__ z,
float* __restrict__ output,
int total_elements
) {
int idx = blockIdx.x * blockDim.x + threadIdx.x;
if (idx < total_elements) {
float z_val = z[idx];
float silu_z = z_val / (1.0f + expf(-z_val));
output[idx] = y[idx] * silu_z;
}
}
torch::Tensor rmsnorm_cuda(torch::Tensor input, torch::Tensor weight) {
TORCH_CHECK(input.is_cuda(), "Input must be a CUDA tensor");
TORCH_CHECK(input.dim() == 3, "Input must be 3D");
int B = input.size(0);
int L = input.size(1);
int D = input.size(2);
auto output = torch::empty_like(input);
int total = B * L;
int threads = 256;
int blocks = (total + threads - 1) / threads;
rmsnorm_kernel<<<blocks, threads>>>(
input.data_ptr<float>(),
weight.data_ptr<float>(),
output.data_ptr<float>(),
B, L, D
);
return output;
}
torch::Tensor selective_scan_cuda(
torch::Tensor x_branch,
torch::Tensor delta,
torch::Tensor B_t,
torch::Tensor C_t,
torch::Tensor A_log,
torch::Tensor D_param
) {
TORCH_CHECK(x_branch.is_cuda(), "x_branch must be a CUDA tensor");
int B = x_branch.size(0);
int L = x_branch.size(1);
int d_inner = x_branch.size(2);
int N = A_log.size(1);
TORCH_CHECK(N <= N_MAX, "d_state must be <= ", N_MAX);
auto output = torch::empty({B, L, d_inner}, x_branch.options());
int threads = 32;
int blocks_x = (d_inner + threads - 1) / threads;
dim3 blocks(blocks_x, B);
selective_scan_kernel<<<blocks, threads>>>(
x_branch.data_ptr<float>(),
delta.data_ptr<float>(),
B_t.data_ptr<float>(),
C_t.data_ptr<float>(),
A_log.data_ptr<float>(),
D_param.data_ptr<float>(),
output.data_ptr<float>(),
B, L, d_inner, N
);
return output;
}
torch::Tensor gate_cuda(torch::Tensor y, torch::Tensor z) {
TORCH_CHECK(y.is_cuda(), "y must be a CUDA tensor");
TORCH_CHECK(z.is_cuda(), "z must be a CUDA tensor");
auto output = torch::empty_like(y);
int total = y.numel();
int threads = 256;
int blocks = (total + threads - 1) / threads;
gate_kernel<<<blocks, threads>>>(
y.data_ptr<float>(),
z.data_ptr<float>(),
output.data_ptr<float>(),
total
);
return output;
}
"""
mamba_ops = load_inline(
name='mamba_ops',
cpp_sources=[],
cuda_sources=[cuda_source],
functions=['rmsnorm_cuda', 'selective_scan_cuda', 'gate_cuda'],
verbose=True
)
class ModelNew(nn.Module):
def __init__(self, d_model: int, d_state: int, d_conv: int, expand: int):
super().__init__()
self.d_model = d_model
self.d_inner = d_model * expand
self.d_state = d_state
self.d_conv = d_conv
self.norm = nn.RMSNorm(d_model)
self.in_proj = nn.Linear(d_model, self.d_inner * 2, bias=False)
self.conv1d = nn.Conv1d(self.d_inner, self.d_inner, d_conv,
padding=d_conv - 1, groups=self.d_inner, bias=True)
self.activation = nn.SiLU()
self.A_log = nn.Parameter(torch.randn(self.d_inner, d_state))
self.proj_B = nn.Linear(self.d_inner, d_state, bias=False)
self.proj_C = nn.Linear(self.d_inner, d_state, bias=False)
self.proj_delta = nn.Linear(self.d_inner, self.d_inner, bias=True)
self.D = nn.Parameter(torch.ones(self.d_inner))
self.out_proj = nn.Linear(self.d_inner, d_model, bias=False)
def forward(self, x: torch.Tensor) -> torch.Tensor:
B, L, D = x.shape
residual = x
# RMSNorm
x = mamba_ops.rmsnorm_cuda(x, self.norm.weight)
# Linear projection and split
xz = self.in_proj(x)
x_branch, z = xz.chunk(2, dim=-1)
x_branch = x_branch.contiguous()
z = z.contiguous()
# Conv1d + SiLU
x_branch = x_branch.transpose(1, 2)
x_branch = self.conv1d(x_branch)[..., :L]
x_branch = x_branch.transpose(1, 2).contiguous()
x_branch = self.activation(x_branch)
# Projections for SSM
delta = torch.nn.functional.softplus(self.proj_delta(x_branch))
B_t = self.proj_B(x_branch)
C_t = self.proj_C(x_branch)
# Selective scan (SSM)
y = mamba_ops.selective_scan_cuda(x_branch, delta, B_t, C_t, self.A_log, self.D)
# Gate with z
y = mamba_ops.gate_cuda(y, z)
# Output projection
y = self.out_proj(y)
# Residual connection
y = y + residual
return y
```
This should be a complete, working implementation. Let me submit it.
call → submit_kernel
callsubmit_kernel
{
"kernel_code": "import torch\nimport torch.nn as nn\nfrom torch.utils.cpp_extension import load_inline\n\ncuda_source = \"\"\"\n#include <torch/extension.h>\n#include <cuda.h>\n#include <cuda_runtime.h>\n#include <cmath>\n\n#define N_MAX 16\n\n__global__ void rmsnorm_kernel(\n const float* __restrict__ input,\n const float* __restrict__ weight,\n float* __restrict__ output,\n int B, int L, int D\n) {\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n int total = B * L;\n \n if (idx < total) {\n float sum_sq = 0.0f;\n for (int d = 0; d < D; d++) {\n float val = input[idx * D + d];\n sum_sq += val * val;\n }\n \n float rms_inv = rsqrtf(sum_sq / D + 1e-5f);\n \n for (int d = 0; d < D; d++) {\n output[idx * D + d] = input[idx * D + d] * rms_inv * weight[d];\n }\n }\n}\n\n__global__ void selective_scan_kernel(\n const float* __restrict__ x_branch,\n const float* __restrict__ delta,\n const float* __restrict__ B_t,\n const float* __restrict__ C_t,\n const float* __restrict__ A_log,\n const float* __restrict__ D_param,\n float* __restrict__ output,\n int B, int L, int d_inner, int N\n) {\n int b = blockIdx.y;\n int c = blockIdx.x * blockDim.x + threadIdx.x;\n \n if (b >= B || c >= d_inner) return;\n \n float A[N_MAX];\n for (int n = 0; n < N; n++) {\n A[n] = -expf(A_log[c * N + n]);\n }\n \n float h[N_MAX];\n for (int n = 0; n < N; n++) {\n h[n] = 0.0f;\n }\n \n float D_val = D_param[c];\n \n int base_x = b * L * d_inner + c;\n int base_BC = b * L * N;\n \n for (int t = 0; t < L; t++) {\n float delta_t = delta[base_x + t * d_inner];\n float x_t = x_branch[base_x + t * d_inner];\n \n const float* B_t_ptr = B_t + base_BC + t * N;\n const float* C_t_ptr = C_t + base_BC + t * N;\n \n float y_t = 0.0f;\n for (int n = 0; n < N; n++) {\n float dA_t = expf(delta_t * A[n]);\n float dB_t = delta_t * B_t_ptr[n];\n h[n] = dA_t * h[n] + x_t * dB_t;\n y_t += h[n] * C_t_ptr[n];\n }\n y_t += D_val * x_t;\n \n output[base_x + t * d_inner] = y_t;\n }\n}\n\n__global__ void gate_kernel(\n const float* __restrict__ y,\n const float* __restrict__ z,\n float* __restrict__ output,\n int total_elements\n) {\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n if (idx < total_elements) {\n float z_val = z[idx];\n float silu_z = z_val / (1.0f + expf(-z_val));\n output[idx] = y[idx] * silu_z;\n }\n}\n\ntorch::Tensor rmsnorm_cuda(torch::Tensor input, torch::Tensor weight) {\n TORCH_CHECK(input.is_cuda(), \"Input must be a CUDA tensor\");\n TORCH_CHECK(input.dim() == 3, \"Input must be 3D\");\n \n int B = input.size(0);\n int L = input.size(1);\n int D = input.size(2);\n \n auto output = torch::empty_like(input);\n \n int total = B * L;\n int threads = 256;\n int blocks = (total + threads - 1) / threads;\n \n rmsnorm_kernel<<<blocks, threads>>>(\n input.data_ptr<float>(),\n weight.data_ptr<float>(),\n output.data_ptr<float>(),\n B, L, D\n );\n \n return output;\n}\n\ntorch::Tensor selective_scan_cuda(\n torch::Tensor x_branch,\n torch::Tensor delta,\n torch::Tensor B_t,\n torch::Tensor C_t,\n torch::Tensor A_log,\n torch::Tensor D_param\n) {\n TORCH_CHECK(x_branch.is_cuda(), \"x_branch must be a CUDA tensor\");\n \n int B = x_branch.size(0);\n int L = x_branch.size(1);\n int d_inner = x_branch.size(2);\n int N = A_log.size(1);\n \n TORCH_CHECK(N <= N_MAX, \"d_state must be <= \", N_MAX);\n \n auto output = torch::empty({B, L, d_inner}, x_branch.options());\n \n int threads = 32;\n int blocks_x = (d_inner + threads - 1) / threads;\n dim3 blocks(blocks_x, B);\n \n selective_scan_kernel<<<blocks, threads>>>(\n x_branch.data_ptr<float>(),\n delta.data_ptr<float>(),\n B_t.data_ptr<float>(),\n C_t.data_ptr<float>(),\n A_log.data_ptr<float>(),\n D_param.data_ptr<float>(),\n output.data_ptr<float>(),\n B, L, d_inner, N\n );\n \n return output;\n}\n\ntorch::Tensor gate_cuda(torch::Tensor y, torch::Tensor z) {\n TORCH_CHECK(y.is_cuda(), \"y must be a CUDA tensor\");\n TORCH_CHECK(z.is_cuda(), \"z must be a CUDA tensor\");\n \n auto output = torch::empty_like(y);\n int total = y.numel();\n \n int threads = 256;\n int blocks = (total + threads - 1) / threads;\n \n gate_kernel<<<blocks, threads>>>(\n y.data_ptr<float>(),\n z.data_ptr<float>(),\n output.data_ptr<float>(),\n total\n );\n \n return output;\n}\n\"\"\"\n\nmamba_ops = load_inline(\n name='mamba_ops',\n cpp_sources=[],\n cuda_sources=[cuda_source],\n functions=['rmsnorm_cuda', 'selective_scan_cuda', 'gate_cuda'],\n verbose=True\n)\n\nclass ModelNew(nn.Module):\n def __init__(self, d_model: int, d_state: int, d_conv: int, expand: int):\n super().__init__()\n self.d_model = d_model\n self.d_inner = d_model * expand\n self.d_state = d_state\n self.d_conv = d_conv\n \n self.norm = nn.RMSNorm(d_model)\n self.in_proj = nn.Linear(d_model, self.d_inner * 2, bias=False)\n self.conv1d = nn.Conv1d(self.d_inner, self.d_inner, d_conv,\n padding=d_conv - 1, groups=self.d_inner, bias=True)\n self.activation = nn.SiLU()\n \n self.A_log = nn.Parameter(torch.randn(self.d_inner, d_state))\n self.proj_B = nn.Linear(self.d_inner, d_state, bias=False)\n self.proj_C = nn.Linear(self.d_inner, d_state, bias=False)\n self.proj_delta = nn.Linear(self.d_inner, self.d_inner, bias=True)\n self.D = nn.Parameter(torch.ones(self.d_inner))\n \n self.out_proj = nn.Linear(self.d_inner, d_model, bias=False)\n \n def forward(self, x: torch.Tensor) -> torch.Tensor:\n B, L, D = x.shape\n residual = x\n \n # RMSNorm\n x = mamba_ops.rmsnorm_cuda(x, self.norm.weight)\n \n # Linear projection and split\n xz = self.in_proj(x)\n x_branch, z = xz.chunk(2, dim=-1)\n x_branch = x_branch.contiguous()\n z = z.contiguous()\n \n # Conv1d + SiLU\n x_branch = x_branch.transpose(1, 2)\n x_branch = self.conv1d(x_branch)[..., :L]\n x_branch = x_branch.transpose(1, 2).contiguous()\n x_branch = self.activation(x_branch)\n \n # Projections for SSM\n delta = torch.nn.functional.softplus(self.proj_delta(x_branch))\n B_t = self.proj_B(x_branch)\n C_t = self.proj_C(x_branch)\n \n # Selective scan (SSM)\n y = mamba_ops.selective_scan_cuda(x_branch, delta, B_t, C_t, self.A_log, self.D)\n \n # Gate with z\n y = mamba_ops.gate_cuda(y, z)\n \n # Output projection\n y = self.out_proj(y)\n \n # Residual connection\n y = y + residual\n \n return y"
}