reasoning
Let me analyze the reference model:
```python
class Model(nn.Module):
def forward(self, q, k, v, gate):
batch_size, seq_len, num_heads, head_dim = q.shape
state = torch.zeros(batch_size, num_heads, head_dim, head_dim, dtype=q.dtype, device=q.device)
outputs = []
for t in range(seq_len):
gate_t = gate[:, t].unsqueeze(-2) # (B, 1, H, D)
kv_outer = k[:, t].unsqueeze(-1) * v[:, t].unsqueeze(-2) # (B, H, D, D)
state = gate_t * state + kv_outer # (B, H, D, D)
y_t = (q[:, t].unsqueeze(-1) * state).sum(dim=-2) # (B, H, D)
outputs.append(y_t)
return torch.stack(outputs, dim=1) # (B, T, H, D)
```
This is a linear recurrence where:
- `state[b, h, i, j]` at step t is updated by: `state[b,h,i,j] = gate[b,t,h,j] * state[b,h,i,j] + k[b,t,h,i] * v[b,t,h,j]`
- Output at step t: `y[b,t,h,j] = sum_i q[b,t,h,i] * state[b,h,i,j]`
Wait, let me be more careful about dimensions:
- `q, k, v`: (B, T, H, D)
- `gate`: (B, T, H, D)
- `gate_t = gate[:, t].unsqueeze(-2)`: (B, 1, H, D)
- `state`: (B, H, D, D)
- `gate_t * state`: gate broadcasts: (B, H, D, D) * gate_t -> (B, 1, H, D) broadcasts to (B, H, D, D)
So `gate_t[b, h, :, j]` is the same for all i, i.e., gate affects each column j of state independently.
- `kv_outer = k[:, t].unsqueeze(-1) * v[:, t].unsqueeze(-2)`: k: (B, H, D, 1), v: (B, H, 1, D), result: (B, H, D, D)
`kv_outer[b, h, i, j] = k[b, t, h, i] * v[b, t, h, j]`
- `state = gate_t * state + kv_outer`
`state[b, h, i, j] = gate[b, t, h, j] * state[b, h, i, j] + k[b, t, h, i] * v[b, t, h, j]`
- `y_t = (q[:, t].unsqueeze(-1) * state).sum(dim=-2)`
`q[:, t].unsqueeze(-1)`: (B, H, D, 1)
`state`: (B, H, D, D)
`q[:, t].unsqueeze(-1) * state`: (B, H, D, D)
sum over dim=-2 (the i dimension, dim=2): sum_i q[b, t, h, i] * state[b, h, i, j]
So: `y[b, t, h, j] = sum_i q[b, t, h, i] * state_t[b, h, i, j]`
And `state_t[b, h, i, j] = gate[b, t-1, h, j] * state_{t-1}[b, h, i, j] + k[b, t-1, h, i] * v[b, t-1, h, j]`
This is like a linear RNN where each channel j has its own decay. But the remarkable thing is that for each (b, h, j), the state is a vector over i, and each element decays by the same gate.
Actually, looking more carefully:
`state[b, h, :, j]` is a vector of size D (indexed by i).
The recurrence for this vector is:
`state_t[b, h, i, j] = gate[b, t-1, h, j] * state_{t-1}[b, h, i, j] + k[b, t-1, h, i] * v[b, t-1, h, j]`
For a fixed b, h, j, let me define:
- `s_t[i] = state_t[b, h, i, j]`
- `g_t = gate[b, t, h, j]`
- `k_t[i] = k[b, t, h, i]`
- `v_t = v[b, t, h, j]`
Then:
`s_t[i] = g_{t-1} * s_{t-1}[i] + k_{t-1}[i] * v_{t-1}`
Output:
`y[b, t, h, j] = sum_i q[b, t, h, i] * s_t[i]`
This can be unrolled:
`s_t[i] = sum_{tau=0}^{t-1} (k[tau][i] * v[tau]) * prod_{p=tau}^{t-1} g_p`
(with the convention that empty product = 1, but wait, the initial state is 0, and at step t we use gate from step t to update state for step t+1 output)
Actually, looking at the loop:
- t=0: state = 0; state = gate[0] * 0 + kv[0]; y[0] = q[0] @ state
- t=1: state = gate[1] * kv[0] + kv[1]; y[1] = q[1] @ state
- t=2: state = gate[2] * (gate[1] * kv[0] + kv[1]) + kv[2] = gate[2]*gate[1]*kv[0] + gate[2]*kv[1] + kv[2]
So for the output at time t, the state is:
`state_t = sum_{tau=0}^{t} kv_tau * prod_{p=tau+1}^{t} gate_p`
where `prod_{p=t+1}^{t} gate_p = 1` (empty product).
So `y[b, t, h, j] = sum_i q[b, t, h, i] * sum_{tau=0}^{t} k[b, tau, h, i] * v[b, tau, h, j] * prod_{p=tau+1}^{t} gate[b, p, h, j]`
= `sum_{tau=0}^{t} v[b, tau, h, j] * prod_{p=tau+1}^{t} gate[b, p, h, j] * sum_i q[b, t, h, i] * k[b, tau, h, i]`
= `sum_{tau=0}^{t} v[b, tau, h, j] * prod_{p=tau+1}^{t} gate[b, p, h, j] * (q[b, t, h] @ k[b, tau, h])`
Wait, but q and k are for the same timestep in standard attention. Here q[t] attends to past k[tau].
Let me define:
- `qk_dot[b, t, tau, h] = sum_i q[b, t, h, i] * k[b, tau, h, i]`
Then `y[b, t, h, j] = sum_{tau=0}^{t} v[b, tau, h, j] * qk_dot[b, t, tau, h] * prod_{p=tau+1}^{t} gate[b, p, h, j]`
This can be computed as a cumulative scan. Let's define for fixed b, h, j:
`y[b, t, h, j] = sum_{tau=0}^{t} qk_dot[b, t, tau, h] * v[b, tau, h, j] * prod_{p=tau+1}^{t} gate[b, p, h, j]`
For a fixed output position t, and fixed channel j, we can think of this as a dot product between:
- `qk_dot[b, t, 0:t, h]`
- `v[b, 0:t, h, j]` with scaling by the gate products.
Actually, one approach to compute this efficiently:
1. Compute `qk_dot` for all (t, tau) pairs - this is O(T^2 * B * H * D). But this materializes an intermediate of size (B, H, T, T).
2. Then for each j, compute weighted sums with gate products.
But we can also compute this iteratively using a scan:
For fixed b, h, j, define `acc_t = sum_{tau=0}^{t} qk_dot[t, tau] * v[tau] * prod_{p=tau+1}^{t} gate_p`
This satisfies:
`acc_t = gate_t * acc_{t-1} + qk_dot[t, t] * v[t]`
Wait, that's not quite right because qk_dot[t, tau] depends on both t and tau.
The issue is that for the standard recurrence, we had:
`s_t[i] = gate_{t-1} * s_{t-1}[i] + k_{t-1}[i] * v_{t-1}`
`y_t[j] = sum_i q_t[i] * s_t[i]`
This is exactly what the loop computes. To vectorize:
For each (b, h, i, j), we have a 1D recurrence over t:
`state_t = gate_{t-1, j} * state_{t-1} + k_{t-1, i} * v_{t-1, j}`
`y_{t, j} = sum_i q_{t, i} * state_t(i, j)`
We can compute this with a parallel scan, but given the dimensions B=2, H=4, T=128, D=24, the sequential loop is actually not that bad in terms of compute. The main issue is that PyTorch's Python loop is slow.
Let me estimate:
- B=2, H=4, T=128, D=24
- Total "state" elements: B * H * D * D = 2 * 4 * 24 * 24 = 4608
- Per step: D*D multiplies for state update, D*D multiplies for output
- Total operations: O(B * H * T * D^2) = 2 * 4 * 128 * 576 ≈ 590k operations
This is tiny. The PyTorch overhead from the Python loop is the main bottleneck. A custom CUDA kernel can easily beat this.
Strategy: Write a CUDA kernel that processes the loop efficiently in a single kernel launch.
Given the small D=24, I can have each thread block handle one (b, h) pair, and threads within the block process different j values.
For each (b, h):
- We have D different recurrences (one for each j), and each recurrence has a D-dimensional state vector.
- Actually, for each j, the state is a vector of size D (indexed by i).
- The recurrence is: `state_t[i] = gate_{t-1}[j] * state_{t-1}[i] + k_{t-1}[i] * v_{t-1}[j]`
- Output: `y_t[j] = sum_i q_t[i] * state_t[i]`
For a fixed j, we need to maintain a vector of D floats.
For each step t:
1. `state[i] = gate[t, j] * state[i] + k[t, i] * v[t, j]` for all i
2. `y[t, j] = sum_i q[t, i] * state[i]`
We can parallelize over i and j within the block. Let's say we use D threads, where each thread handles one j. But then each thread needs to process all D values of i sequentially - that's 24 operations per step, times 128 steps = 3072 per thread. That's also fine.
Actually, maybe better: each thread can handle one (i, j) pair. We have D*D = 576 pairs. With 576 threads, we can process all of them in parallel. But then we need to do a reduction for output.
Or: We can store the state in shared memory. Each of D threads handles one column j of the state (size D). Each thread iterates over t and for each t, updates all D elements in its column, then computes dot(q_t, state_column_j) and stores to output.
Wait, if each thread handles one j:
- State[j][i] stored in registers or shared memory. But we need D floats per thread = 24 floats = 96 bytes. This fits easily in registers.
- For each step t:
- `state[i] = gate[t, j] * state[i] + k[t, i] * v[t, j]` for i=0..D-1
- `dot = sum_i q[t, i] * state[i]`
- `y[t, j] = dot`
This is very natural. We launch a grid of (B, H) blocks, each with D threads. Total blocks = B*H = 8, total threads = 8 * 24 = 192.
But blocks are few. Maybe we need to use more parallelism. Since T=128 is small, we could have each block process some steps, but step t depends on step t-1, so we can't parallelize over t.
Alternative: process multiple (b, h) pairs? B*H = 8. Maybe interleave multiple j values per thread?
Actually with H100 we have many SMs. 8 blocks is very few. We might want to launch more threads per block to saturate the GPU, or have each block handle a chunk of the problem differently.
Wait, we can also think of this as having each block handle one (b, h, j) with multiple threads processing the D elements. But then we'd have B*H*D = 2*4*24 = 192 blocks. That's still not huge but better.
Or we can use more threads per block. For instance, block_size = min(D * something). Let's say 256 threads per block. Each block handles floor(256/D) = 10 j values, or we can have each block handle one (b, h) and the 256 threads can process D values with some threads idle, or we can have each thread process multiple j's.
Let me try a simple approach first and see if it compiles and is correct. Then we can optimize.
Each (b, h) has D independent recurrences (one per j). Each recurrence updates a D-dimensional vector.
Kernel design:
- Grid: (B * H) blocks
- Block: some number of threads, say D threads (one per j)
- Each thread j maintains float state[D] in registers
- Loop over t from 0 to T-1:
- gate_j = gate[flat_index + j]
- v_j = v[flat_index + j]
- k_i = k[...] for i from 0 to D-1
- state_i = gate_j * state_i + k_i * v_j
- dot product with q_i
- write output
But we need to read k and q for all i. If each thread only has one j, then each thread needs to read k[all i] and q[all i]. These loads aren't coalesced well if each thread reads its own j and loops over i.
Let's think about memory access patterns:
- q, k, v: (B, T, H, D) or could be stored as (B, H, T, D) for better access. The inputs are (B, T, H, D) though. We have to work with that.
- gate: (B, T, H, D)
For coalescing: adjacent threads should access adjacent memory. So thread j should access element j.
If block has D threads, and thread j handles recurrence j:
- Read gate[t, j]: adjacent threads read adjacent j -> coalesced
- Read v[t, j]: coalesced
- Read k[t, i]: this is tricky. If thread j loops over i, all threads read the same k[t, i] for a given i - broadcast load, not coalesced in the j dimension but it can be loaded via L1 broadcast.
- Similarly for q[t, i].
Actually, k[t, :] and q[t, :] are the same across all j in a block. So we can load them into shared memory once per block.
Let me reconsider: each block processes one (b, h). We load k_t and q_t into shared memory. Then each thread j updates its state and computes the dot product.
But if each thread handles one j and processes D elements of state, each thread needs state[D] registers. D=24. That's fine.
Wait, an even better idea: what if each thread handles one i and processes all j? Then the recurrence for state[i][j] is:
`state[i][j] = gate[j] * state[i][j] + k[i] * v[j]`
This can be vectorized! The operation is essentially:
`state = gate * state + k[i] * v` where gate and v are vectors of size D, and k[i] is a scalar. This is a scaled vector addition plus element-wise multiplication.
If each thread handles one i, then it needs to store state[j] for j=0..D-1 = 24 floats. That's also fine. But then the output `y[j] = sum_i q[i] * state[i][j]` requires reduction over i.
For each step t:
- Thread i computes: `state[j] = gate[j] * state[j] + k[i] * v[j]` for all j=0..D-1
- After all threads do this, we need to compute `y[j] = sum_i q[i] * state[i][j]`. This requires a reduction across threads.
If we have D threads (one per i), after updating state, each thread has its own row state[i][j]. To compute y[j], we need `sum_i q[i] * state[i][j]`. We can transpose and reduce.
Actually, if thread i stores state[j] for all j, then to compute y[j] we'd need to access state[i][j] from all threads i, multiply by q[i], and sum. This is a bit more complex because of cross-thread communication.
Alternative: use warp-level operations. With 24 threads, we could use a warp. But warps are 32 threads, so we'd have 8 threads inactive.
Let me go back to: each thread handles one j. Then:
- Shared memory: q_t[D], k_t[D]
- Thread j loads gate_t[j], v_t[j]
- For i=0..D-1: state_i = gate_t[j] * state_i + k_t[i] * v_t[j]
- dot = sum_i q_t[i] * state_i
- write y_t[j] = dot
This is very straightforward. The issue is that it only uses D threads per block.
To increase occupancy, we can have each block handle multiple (b, h) pairs, but B*H=8 total, so we'd have at most 8 blocks.
We could also increase threads by having each thread handle multiple j values. Instead of 1 thread per j, we can process M j values per thread. Block size = D/M. To maximize occupancy with 8 blocks, maybe we should make each block bigger.
But wait - H100 has many SMs and handles small kernel launches fine. The main cost is just loading memory. Let me calculate memory traffic:
- Input reads: q, k, v, gate: each is B*T*H*D floats = 2*128*4*24 = 24,576 floats = ~100KB
- Output writes: B*T*H*D floats = ~100KB
- Total memory: ~200KB
This is tiny! The kernel will be memory-bandwidth bound with very little compute. The reference implementation in Python is slow because of Python loop overhead (128 iterations), not because of computation.
With a custom kernel, we just need to eliminate the Python loop and do all the work in CUDA. Even a simple sequential kernel in CUDA will be much faster because:
1. No Python overhead
2. Single kernel launch
3. Data stays in registers/SharedMem between steps
So the plan is: write a kernel where each block computes the output for one (b, h), and each thread processes one j (or some subset of j's). Use shared memory for q_t and k_t to avoid repeated loads.
Actually, let me optimize further. With D=24, we can fit the entire state in shared memory: B*H blocks, each with D*D floats = 576 * 4 = 2304 bytes. A block could load the entire state in shared memory.
But for a simple first try:
Kernel signature: `(const float* q, const float* k, const float* v, const float* gate, float* out, int B, int T, int H, int D)`
Where the arrays are in (B, T, H, D) layout. Index: `((b * T + t) * H + h) * D + d`.
For each block (b, h):
- Thread j (0 <= j < D) initialize state[j][i] = 0 for i=0..D-1 (in registers as `float state[24]` or dynamically). But dynamic arrays in thread local would be on stack. Better to use fixed-size if D<=24, but D is runtime parameter. Actually in the problem, D=24 is fixed but the kernel should accept it as a parameter. We can use a compile-time constant or just allocate a fixed max size.
Wait, we don't know the max D at compile time from the signature, but we can make it a template parameter or just use a shared memory array for state. But the state per thread is D floats. We can put it in shared memory: `__shared__ float state[BLOCK_DIM][MAX_D]` where BLOCK_DIM is block size. Since D=24 and block size could be D=24, that's 24*24*4 = 2304 bytes. Very manageable.
Actually, wait. The original state is (B, H, D, D). For each block (b, h), the state is D x D. That maps nicely to shared memory.
But with the recurrence, the state is updated as: `state[i][j] = gate[j] * state[i][j] + k[i] * v[j]`.
If we store state in shared memory as state[i][j], thread (i,j) can update state[i][j]. But i and j range over D x D = 576. So we could launch a 2D thread block of (D, D) = 24 x 24 = 576 threads. But then computing the output `y[j] = sum_i q[i] * state[i][j]` requires reduction over i, which means threads with the same j need to reduce.
Alternatively, we could process sequentially but keep state in shared memory. For each step t, we load k_t, q_t, gate_t into shared memory. Then each thread in a D x D block updates one element of state. After that, we do a reduction for each j across i.
But a 2D block has 576 threads, and we'd have B*H = 8 blocks. Total 4608 threads. H100 can run many more, but this is fine.
However, we only need D^2 operations per step for state update, and D^2 operations for output. With 576 threads, that's 1 op per thread per step. Over 128 steps, very little work. But again, the main win is eliminating Python overhead.
Let me try the simpler 1D block approach with each thread handling one j, and keeping state in registers.
For thread j in block (b, h):
```cuda
float state[24]; // or dynamic: if D <= 32 can use small array
for (int i = 0; i < D; i++) state[i] = 0.0f;
for (int t = 0; t < T; t++) {
float gate_j = gate[((b * T + t) * H + h) * D + j];
float v_j = v[((b * T + t) * H + h) * D + j];
for (int i = 0; i < D; i++) {
float k_i = k[((b * T + t) * H + h) * D + i];
state[i] = gate_j * state[i] + k_i * v_j;
}
float dot = 0.0f;
for (int i = 0; i < D; i++) {
float q_i = q[((b * T + t) * H + h) * D + i];
dot += q_i * state[i];
}
out[((b * T + t) * H + h) * D + j] = dot;
}
```
This is clean and correct. Let's think about performance.
- For each (b, h) and each t, thread j reads gate_j and v_j once. There are D threads, so all gate and v values for this (b, h, t) are read exactly once, coalesced.
- Each thread reads all k[i] for i=0..D-1. Since all threads read the same k values, this is broadcast from L1 cache. Actually each thread reads the entire k_t array of D floats. Total k reads per step per block = D threads * D floats = D^2 floats. But k_t has only D floats. So we read D times more k than necessary. Same for q.
To fix this, we can load k_t and q_t into shared memory once per block:
```cuda
__shared__ float k_shared[24];
__shared__ float q_shared[24];
__shared__ float gate_shared[24];
__shared__ float v_shared[24];
// thread 0 loads, or better, all threads load their own j then sync
// Actually for coalescing, thread j loads gate[j], v[j], and we need k[i] and q[i] for all i.
// Thread j can load k[j], q[j]. Then __syncthreads().
// Then all threads can access k_shared[i], q_shared[i].
```
This is much better:
- Thread j loads: gate[j], v[j], k[j], q[j] - all coalesced
- __syncthreads()
- Thread j computes using shared memory
Total shared memory per block = 4 * D * sizeof(float) = 4 * 24 * 4 = 384 bytes! Negligible.
But wait, block has only D=24 threads. To improve occupancy, maybe we should handle multiple (b, h) pairs per block or have more threads. But let's see - with 8 blocks of 24 threads, it's 192 threads total. H100 can run thousands of threads concurrently. However, each thread does 128 * (24 + 24 + 24) ≈ ~9000 FLOPs. 192 threads * 9000 FLOPs = 1.7M FLOPs total. At 67 TFLOPS, this is 0.000025 ms. Memory bandwidth bound at 3.35 TB/s for 200KB is 0.00006 ms.
Actually the arithmetic is so small that this will basically run instantly if memory bound, or take a few microseconds. The Python loop version probably takes milliseconds due to Python overhead.
So even 8 blocks of 24 threads should be fine. But maybe PyTorch launch overhead dominates. We want the kernel to be as fast as possible.
Let me implement a clean version first.
One issue: dynamic array `float state[24]` only works if we know the size at compile time. Since D is passed at runtime, we might need to either:
1. Use a maximum size: `float state[MAX_D]` where MAX_D = 64 or something. But D can be anything.
2. Allocate dynamically or use shared memory for state.
Actually in C++ CUDA, you can do `float state[24]` as a fixed-size array. The inputs specify D=24 at runtime, and we can just iterate up to D. If D > 24, we'd need a larger array, but for this problem D=24. To be safe for any reasonable D, we can set a compile-time maximum.
Or we can put the state array in shared memory: `__shared__ float smem_state[BLOCK_SIZE][MAX_D];` where BLOCK_SIZE = D. But then each thread accesses its own row, which is fine.
Actually, keeping state in registers is better for performance. Let's use template parameter or just a fixed max. I'll define `const int MAX_D = 64;` and use `float state[MAX_D];`.
Wait, I need to be careful. CUDA registers spill to local memory if we use too many. But `float state[24]` is 24 registers for the array + some index registers. That's fine.
Let me write the kernel.
Another consideration: can we further optimize? There are T steps, each step needs:
- Read 4*D floats
- Compute state update: D multiplies + D fmas = 2D ops
- Compute output: D multiplies + D-1 adds = 2D ops
- Write D floats
Per step: 4D loads + D stores + ~4D ops
Total: 5D*T memory accesses per (b,h)
Total FLOPs: 4D*T per (b,h)
Arithmetic intensity: (4*D*T) / (5*D*T*4 bytes) = 4 / 20 = 0.2 FLOPs/byte
At 3.35 TB/s bandwidth, we'd get 0.2 * 3.35e12 = 670 GFLOPS. Our compute is tiny, so memory bound.
Can we reduce memory footprint? The outputs are the same size as inputs. The state is computed on-the-fly.
We could potentially fuse the computation so we don't write intermediate state, but we must write outputs.
Actually, looking at the total data: inputs q, k, v, gate are 4 * 100KB = 400KB. Output is 100KB. Total 500KB.
At 3.35 TB/s: 500KB / 3.35TB/s = 0.000149 ms = 149 ns. That's the absolute minimum. In practice, kernel launches take ~5-10 microseconds.
Wait, this is way too optimistic because we're doing random accesses and the cache might not be perfect. But even 10-20 microseconds would be a huge speedup over the Python loop version (which might take 1-5 ms).
Let me check if there's a more efficient algorithm.
Alternative formulation: precompute cumulative products of gates.
For each j: `G[tau, t] = prod_{p=tau+1}^{t} gate_p`
Then output j at time t: `y[t, j] = sum_{tau=0}^{t} v[tau, j] * (q[t] @ k[tau]) * G[tau, t]`
This requires an O(T^2) computation for each j. With T=128, D=24, B=2, H=4:
- qk_dot: B*H*T*T*D = 2*4*128*128*24 = 3.1M ops
- Weighted sum: B*H*T*T*D = 3.1M ops
Total ~6M ops. But materializing intermediate of size B*H*T*T = 2*4*128*128 = 131K floats = 0.5MB. This is not bad, but the sequential scan is only ~600K ops and doesn't require O(T^2) memory.
The O(T^2) approach doesn't help because our sequential scan is already memory bandwidth bound and the O(T^2) version does more compute with worse memory access.
So stick with the sequential scan kernel.
To make the kernel even faster, we can consider:
1. Increasing occupancy by handling multiple (b, h) per block or using more threads
2. Unrolling the inner loops
3. Using vectorized loads/stores
Since D=24 is divisible by 4, we can load/store 4 floats at a time using float4. This could help with memory bandwidth.
But let's start simple.
Actually, I just realized: the kernel will be launched with 8 blocks (B*H). On H100, to get good performance, we want more blocks to fill all SMs. But since the total work is so small, it's not about SMs - it's about latency. More blocks would help hide latency if there were memory stalls, but with only 8 blocks, we'd underutilize the GPU. However, each block has 24 threads, so total 192 threads. GPU can run many more, but that's OK because the work is tiny.
To increase parallelism, we could process multiple time steps in parallel - but they must be sequential due to recurrence. Unless we do a parallel scan (Blelloch), but for 128 steps that's overkill.
Wait! There's another approach: since D=24 is small, we can have each SM process the entire sequence for one (b, h). Multiple SMs can process different (b, h). But we only have 8 (b, h) pairs. So at most 8 SMs are used. With 132 SMs on H100, we're using <10%.
Can we parallelize across another dimension? We can process each (b, h, j) independently if we think about it differently. In fact, for each j, the recurrence is independent! And there are B*H*D = 2*4*24 = 192 independent recurrences. Each recurrence updates a D-dimensional vector over T steps.
In the original formulation, for fixed (b, h, j), state is a D-dimensional vector (over i). But if we allocate one block per (b, h, j), then:
- Each block has some threads, say D threads.
- We have 192 blocks total.
- Each block processes T steps updating a D-dimensional state.
- But then at each step, thread i updates `state[i] = gate * state[i] + k[i] * v`. Thread i needs k[i] and updates state[i]. Then we need to compute dot product of q and state for output.
- All threads in a block need to read q, k, v, gate for each step.
- With D threads per block and 192 blocks, total threads = 192*24 = 4608. Better utilization!
Actually, if each block handles one (b, h, j), then:
- Threads 0..D-1 in the block process i=0..D-1
- Shared memory: we can load k_t[0..D-1] into shared memory, q_t[0..D-1] into shared memory, gate_t (same for all threads), v_t (same for all threads).
- state[i] is stored in each thread's registers or shared memory. If thread i holds state[i], then state is naturally distributed.
- For each step t:
- Load gate_t, v_t (into shared memory or broadcast)
- Thread i loads k_t[i] into shared memory or keeps in register
- __syncthreads()
- Thread i: state_i = gate_t * state_i + k_t[i] * v_t
- Write state_i to shared memory (if not already there)
- __syncthreads()
- Thread i computes partial: q_t[i] * state_i
- Warp-level or block-level reduction to get dot product
- Write dot to output[j]
Wait, but thread i only has state_i. To compute dot = sum_i q_i * state_i, each thread computes q_i * state_i, then all threads reduce. With D=24, we can do a warp shuffle reduction or tree reduction in shared memory.
But we only need one output scalar per block per step. That's very little parallelism within the block for the reduction. But D=24 is small.
Actually, the output at step t for this (b, h, j) is just one value. So each block writes one float per step. Total output elements per block: T = 128. Total blocks: 192.
Thread i in block (b, h, j):
- Holds state[i] in register
- For t=0..T-1:
- Read k_t[i] (coalesced across i)
- Read q_t[i] (coalesced across i)
- Read gate_t[0] (first thread, or all threads could read and broadcast)
- Read v_t[0] (same)
- state_i = gate * state_i + k_i * v
- __syncthreads() or just do reduction
- partial = q_i * state_i
- reduction over all threads -> dot
- write dot to out[t, j]
This seems more complex due to the reduction, and memory access pattern is similar. But now we have 192 blocks instead of 8, so better SM utilization.
However, the first approach (1 block per (b,h), 1 thread per j) is:
- 8 blocks
- 24 threads per block
- Each thread does 128*(24+24) ops
- No reduction needed
- Output writes are perfectly coalesced (thread j writes out[t, j])
In the second approach (1 block per (b,h,j), 1 thread per i):
- 192 blocks
- 24 threads per block
- Each thread does 128*2 ops
- Need reduction per step
- Output writes: block (b,h,j) writes one value per step. For fixed t, blocks with same (b,h) write different j values. Are these coalesced? The output is (B, T, H, D). For fixed b, t, h, the D values are contiguous. If 24 blocks (one per j) each write to out[t, h, j] simultaneously, each block writes one float. These are not coalesced because each block writes to a different address, and writes are done by potentially different SMs. Actually, coalescing is per warp/block. Each block writes one float using one thread, which is terrible for memory.
In approach 1: all D threads in a block write D contiguous floats simultaneously. This is perfectly coalesced (1 transaction of D*4 bytes).
So approach 1 is better for output writes. Approach 2 is better for block count.
What if we do: each block handles one (b, h). But instead of D threads per block, we use more threads. For example, 256 threads per block. We can process multiple j values per thread.
If block has 256 threads and D=24:
- 256/24 = 10 j values with 16 threads left over? Actually 24 doesn't divide 256.
- Or we can have each thread handle ceil(256/24) = 11 j values? No, 256 threads for D=24 columns means each thread handles one column and we have 256-24 = 232 threads idle. That doesn't help.
What if we process multiple (b, h) pairs per block? 256 threads can process 10 (b, h) pairs with 24 threads each, plus 16 threads idle. 256 = 10*24 + 16. But B*H = 8, so we'd fit all 8 pairs in one block, using 192 threads, with 64 threads idle.
Actually, that's an interesting idea! Put all (b, h) pairs in a single block:
- 192 threads total
- Thread layout: for (b, h) pair idx = thread_id / 24, j = thread_id % 24
- We process all 8 (b, h) pairs simultaneously!
- But then for each step t, we need to load q, k, v, gate for all 8 (b, h) pairs.
- This is: 8 * 24 floats for q, 8 * 24 for k, 8 * 24 for v, 8 * 24 for gate.
- Total per step: 8 * 96 = 768 floats = 3KB per step. Over 128 steps: 384KB.
With 192 threads in one block, we'd load these values sequentially (coalesced within each sub-group). But since there are 8 independent sub-groups, the memory access is spread out.
Actually, we can think of it as 8 blocks fused into one. Each sub-group of 24 threads handles one (b, h). The sub-groups are independent. We can use per-sub-group shared memory or just let each thread load its own values.
Let's think: 192 threads in a block.
- Thread layout: `int bh = threadIdx.x / D; int j = threadIdx.x % D;`
- For each t from 0 to T-1:
- All threads load their `gate[bh, t, j]` and `v[bh, t, j]` (coalesced globally)
- For each sub-group, all threads also need `k[bh, t, i]` for all i. But thread j in sub-group bh only has gate and v. It can load k[bh, t, j] into shared memory. With 24 threads per sub-group, we get all k values. But we need them in shared memory accessible to all threads in the sub-group.
We'd need shared memory layout:
`float k_smem[8][24]; float q_smem[8][24];`
// Actually, each thread j loads k[bh, t, j] into `k_smem[bh][j]` and q into `q_smem[bh][j]`
Then `__syncthreads()` among all 192 threads.
Then each thread can compute its state and output.
This uses shared memory: 8 * 24 * 2 * 4 bytes = 1536 bytes. Very small.
With 192 threads in one block, we only launch 1 block! That means only 1 SM is used out of 132. That's terrible. Even though all 192 threads do independent work, they're on one SM. So we'd be limited by 1 SM's throughput.
So that's worse than 8 blocks on 8 SMs.
What about 2 blocks with 96 threads each, handling 4 (b, h) pairs each? Launch 2 blocks. Or 4 blocks with 48 threads each (2 pairs per block), or 8 blocks with 24 threads each.
8 blocks on 8 SMs is the best among these.
Actually, could we increase threads per block while keeping 8 blocks? For instance, each block handles one (b, h) but has 128 threads. Then each thread handles D/128? No, D=24. So some threads would be idle. Or we could interleave: threads handle time steps? But time steps are sequential.
Wait, what if we use a different algorithm that parallelizes over time? There's a technique called "parallel scan" for linear recurrences. For each j, the recurrence is:
`s_t = g_t * s_{t-1} + k_t * v_t`
This is a standard scan operation. For D=24 channels simultaneously (each state is a vector). But parallel scan has overhead and for T=128 may not be worth it.
Alternatively, we can use a "chunked" approach where each thread block handles a chunk of time steps, computes prefix states, and then combines them. This is the standard parallel-prefix approach for linear recurrences. But with T=128 and the overhead, it's probably not worth it for this problem.
Actually, let me reconsider the compute vs memory. Even with 8 blocks, each block runs on one SM and does ~2000 ops for its entire workload. At 67 TFLOPS per SM, that's ~0.03 microseconds. The memory footprint is tiny and will mostly hit L2 cache. So the kernel should finish very fast. The main overhead is kernel launch (~5-10 microseconds).
But the reference PyTorch implementation does 128 Python loop iterations, each creating new tensors, doing tensor ops, appending to list, and then stacking. This probably takes 0.5-2 milliseconds. So even a simple kernel that takes 10 microseconds will be a 50-200x speedup.
So let's go with the simple, clean approach.
Let me write the code.
```python
import torch
from torch.utils.cpp_extension import load_inline
cuda_source = """
#include <torch/extension.h>
#include <cuda_runtime.h>
#define MAX_D 128
__global__ void gated_delta_attention_kernel(
const float* q,
const float* k,
const float* v,
const float* gate,
float* out,
int B,
int T,
int H,
int D
) {
int bh = blockIdx.x;
int b = bh / H;
int h = bh % H;
int j = threadIdx.x;
if (j >= D) return;
// Local state: D floats per thread
float state[MAX_D];
#pragma unroll
for (int i = 0; i < D; i++) {
state[i] = 0.0f;
}
__shared__ float k_shared[MAX_D];
__shared__ float q_shared[MAX_D];
for (int t = 0; t < T; t++) {
int base = ((b * T + t) * H + h) * D;
// Coalesced load of k and q
k_shared[j] = k[base + j];
q_shared[j] = q[base + j];
float gate_j = gate[base + j];
float v_j = v[base + j];
__syncthreads();
#pragma unroll
for (int i = 0; i < D; i++) {
state[i] = gate_j * state[i] + k_shared[i] * v_j;
}
float dot = 0.0f;
#pragma unroll
for (int i = 0; i < D; i++) {
dot += q_shared[i] * state[i];
}
out[base + j] = dot;
}
}
torch::Tensor gated_delta_attention_cuda(
torch::Tensor q,
torch::Tensor k,
torch::Tensor v,
torch::Tensor gate
) {
int B = q.size(0);
int T = q.size(1);
int H = q.size(2);
int D = q.size(3);
auto out = torch::zeros_like(q);
int num_blocks = B * H;
int num_threads = D;
gated_delta_attention_kernel<<<num_blocks, num_threads>>>(
q.data_ptr<float>(),
k.data_ptr<float>(),
v.data_ptr<float>(),
gate.data_ptr<float>(),
out.data_ptr<float>(),
B, T, H, D
);
return out;
}
"""
# But wait, we need to generate bindings. The functions argument tells load_inline what to expose.
# The wrapper function is `gated_delta_attention_cuda`
```
Wait, there are a few problems with this:
1. `MAX_D = 128` and `float state[128]` per thread uses more registers but is OK. However, for D < 24, this is wasteful but fine.
2. `__syncthreads()` inside a loop can be expensive but with small T it's fine.
3. The `#pragma unroll` with `i < D` where D is not known at compile time won't unroll. But we can have an inner loop with a runtime bound. Actually `D` is passed as a parameter. To get unrolling, we'd need a template parameter or just accept it won't unroll.
Actually, for better performance with D=24, we could template the kernel on D and dispatch. Or we could just not worry about it because 24 iterations is small.
But there's another issue: with 24 threads per block, and we do `__syncthreads()` inside the loop over T, this is a valid warp-level synchronization. But if we have exactly 24 threads in a block of size 24, `__syncthreads()` works fine.
One more thing: if D > warpSize (32), we'd have multiple warps. D=24 < 32, so each block is a single warp. That means `__syncthreads()` acts like a warp sync. But not all threads in the warp are active? If D=24, then threads 0-23 are active and 24-31 are inactive. In CUDA, if some threads in a block are inactive due to `return`, `__syncthreads()` can deadlock! The rule is: all threads in a block must execute `__syncthreads()`. But we have `if (j >= D) return;` at the beginning. If block size is 32 and D=24, threads 24-31 will return and never reach `__syncthreads()`, causing deadlock!
This is critical. To avoid this, we should set block size exactly to D, or ensure block size <= D and no early returns. But block size must be set at launch time. We can set `num_threads = D` and then `if (j >= D)` won't trigger because threadIdx.x ranges from 0 to D-1.
Actually if we launch with `num_threads = D`, then j = threadIdx.x is always < D. So the `if (j >= D) return;` is unnecessary and we can remove it.
Wait, D=24, we launch 24 threads per block. All threads 0..23 are active. No thread returns. Perfect.
But what if someone uses D=25? Then we'd launch 25 threads. CUDA block sizes don't have to be power of 2. But performance might be slightly worse for non-multiples of warp size. But D=25 < 32 means one incomplete warp, which is fine.
Actually, if D is say 100, then we'd have 100 threads = 4 warps (96) + 4 threads in the 5th warp. The last warp has 4 active threads and 28 inactive. That's fine - no `__syncthreads()` issues because there is no return statement.
So we should remove the `if (j >= D) return;` since block size exactly equals D.
Let me rewrite the kernel cleanly:
```cpp
template <int D>
__global__ void gated_delta_attention_kernel(
const float* __restrict__ q,
const float* __restrict__ k,
const float* __restrict__ v,
const float* __restrict__ gate,
float* __restrict__ out,
int B,
int T,
int H
) {
const int bh = blockIdx.x;
const int b = bh / H;
const int h = bh % H;
const int j = threadIdx.x;
float state[D];
for (int i = 0; i < D; i++) {
state[i] = 0.0f;
}
__shared__ float k_smem[D];
__shared__ float q_smem[D];
for (int t = 0; t < T; t++) {
const int base = ((b * T + t) * H + h) * D;
k_smem[j] = k[base + j];
q_smem[j] = q[base + j];
const float gate_j = gate[base + j];
const float v_j = v[base + j];
__syncthreads();
for (int i = 0; i < D; i++) {
state[i] = gate_j * state[i] + k_smem[i] * v_j;
}
float dot = 0.0f;
for (int i = 0; i < D; i++) {
dot += q_smem[i] * state[i];
}
out[base + j] = dot;
}
}
```
Wait, `float state[D]` where D is a template parameter. That's valid C++ and CUDA. But for a function template, we need to instantiate it. And we need to dispatch based on runtime D.
Since D=24 at runtime, and it's specified at module load time by the test, we can either:
1. Write a templated version and instantiate it for common values
2. Use dynamic array or shared memory for state
If we use dynamic array (VLAs aren't standard in C++, but CUDA doesn't support VLAs in device code). We can't do `float state[D}` where D is a runtime variable. That's a VLA and not supported in standard C++ or CUDA.
Options:
1. Use a compile-time constant MAX_D and template on that: `float state[MAX_D]`. Works for D <= MAX_D. This is what I originally thought.
2. Use shared memory for state: `extern __shared__ float smem[];` and allocate per thread. But we'd need smem size = blockSize * D. If blockSize = D, then D*D per block.
3. Use local memory (stack): float* state = new float[D]; but dynamic allocation in device code is tricky.
The simplest is option 1: `float state[MAX_D];` where MAX_D is a compile-time constant. Since D <= 128 typically, we can set MAX_D = 128. Each thread uses up to 128 floats = 512 bytes = 128 registers. This is OK.
But for better performance, maybe template based on D with a switch statement. Let me see:
```cpp
void dispatch_kernel(float* q, ..., int D, ...) {
switch(D) {
case 8: kernel<8><<<...>>>(...); break;
case 16: kernel<16><<<...>>>(...); break;
case 24: kernel<24><<<...>>>(...); break;
case 32: kernel<32><<<...>>>(...); break;
case 64: kernel<64><<<...>>>(...); break;
case 128: kernel<128><<<...>>>(...); break;
default: kernel_general<<<...>>>(...);
}
}
```
This is clean and gives the compiler full unrolling ability. For the problem, D=24. So we can just instantiate for D=24.
But wait, in the switch statement, D is a runtime variable. We can check it and dispatch. And if it's not a known size, fall back to a general kernel.
The general kernel can use `float state[128]` (or whatever max).
Let me write a complete version.
Actually, looking at the problem, `head_dim = 24`. So D is always 24 in tests. But the kernel should handle arbitrary D to be safe. Let me dispatch for D=24 specifically and have a fallback.
Wait, the kernel with `float state[MAX_D]` where MAX_D is a compile-time constant is perfectly fine and avoids template bloat. With MAX_D=128 and D only up to 128, each thread uses at most 128 registers for state. For D=24, that's only 24 registers used.
Let me reconsider whether we need to maximize occupancy. We have 8 blocks of 24 threads each. With 24 threads per block and no `__syncthreads()` divergence issues, each block is one warp. On H100, each SM can run up to 32 blocks simultaneously. So 8 blocks could run on a single SM or be spread across SMs. With only 8 blocks, we're far from saturating the GPU. But the workload is tiny anyway.
To improve occupancy, we could split each (b, h) pair across multiple blocks in the time dimension. For example, each block handles a chunk of T. Since the recurrence is sequential, we'd need prefix sums or a two-pass algorithm. But for T=128, it's not worth it.
What if we process multiple (b, h, j) triples per thread? This increases the work per thread but doesn't increase block count. Since blocks are already spread across SMs, increasing work per thread might just increase latency without improving throughput.
But wait, what if we use a 2D block where threads are (D, something)? If we set blockDim.y = some_factor, we can have block size = D * factor. Each thread handles one j but factor different (b,h) pairs? No, we have only 8 (b,h) pairs.
Actually, we can have each thread handle a portion of the j dimension. For instance, block size = 256, and each thread handles ceil(D/256) = 1 j value... no, 256 > 24.
OR, we can increase block count by processing different ranges of j in different blocks while keeping (b,h) the same. For example, we could have 4 blocks for each (b,h), each handling D/4 = 6 j values. Then total blocks = 8 * 4 = 32. Block size = 6 threads. This is worse because small blocks have overhead.
Hmm, let me just go with the simple version and test it. I can optimize later if needed.
Wait, I should also consider using `__ldg` (load through read-only cache/texture) for q, k, v, gate to potentially improve cache behavior. Or just use `const float* __restrict__` which helps the compiler know there are no aliasing issues.
Also, for the dot product reduction, since D=24, we could use `#pragma unroll` for the reduction if we template. But in the simple version without template, we can still partially unroll by having a switch or just accept 24 iterations.
Actually, let me think about instruction-level parallelism. In the state update loop:
```cpp
for (int i = 0; i < D; i++) {
state[i] = gate_j * state[i] + k_smem[i] * v_j;
}
```
This is a loop-carried dependency only on `state[i]`, but different iterations are over different `state[i]` elements, so they're independent. The compiler can unroll and schedule them efficiently.
Similarly for the dot product.
Let me write the kernel code now. I'll use the simple approach with `state[128]`, and since D is always 24 in practice, we could also add a specialized path for D=24. But let's keep it general first.
Actually, I just realized - if we put state in registers as a fixed-size array but only use the first D elements, the compiler is smart enough to not allocate more registers than needed for the active elements.
One more optimization: we can process the gate and v loads outside of the k_smem, q_smem loads, and since gate_j and v_j are the same for all threads in the block... wait no, gate_j depends on j! Each j has its own gate value. In the original code, gate_t = gate[:, t].unsqueeze(-2) which broadcasts across the D dimension. Wait, let me re-read:
```python
gate_t = gate[:, t].unsqueeze(-2)
state = gate_t * state + kv_outer
```
`gate_t` shape: (B, 1, H, D) - yes, it broadcasts over the i dimension (dim=-2 in state) but is different for each j (dim=-1). So for each j, `gate_j` is a scalar that scales the entire j-th column of state.
So each thread j has its own `gate_j`. That's what we have.
Similarly, v_j is per j.
OK, the algorithm is correct.
One more thought: `k_smem[i]` is accessed by all threads. Is there bank conflict? Each thread reads the same `k_smem[i]` in the i-loop iteration. When all threads in a warp read the same shared memory address, it's a broadcast, which is very efficient (no bank conflicts).
For `q_smem[i]`, same thing - broadcast in the dot product loop.
For writing `out[base + j]`, all threads write adjacent memory, perfectly coalesced.
For loading `k_smem[j] = k[base + j]`, threads write to adjacent shared memory.
For loading `q_smem[j] = q[base + j]`, same.
For loading `gate_j = gate[base + j]` and `v_j = v[base + j]`, adjacent threads read adjacent global memory, perfectly coalesced.
This is a very well-structured kernel in terms of memory access.
One optimization: since the block is exactly 24 threads = one incomplete warp, memory transactions might be padded to 32 elements. We're reading 24 contiguous floats. That should be done in one 128-byte transaction (32 floats) with some wasted bandwidth. But it's fine.
Could we pad to 32 to avoid waste? Not really, because D=24 is fixed.
Actually wait - we don't need to use a block size of exactly D. We could use 32 threads with 8 of them doing nothing. But then k_smem[32], q_smem[32], and state[24] for active threads. But the 8 extra threads do nothing and waste some resources. Block size of 24 is fine.
Let me also think about whether we should prefetch the next time step's q, k, v, gate while computing the current step. But with shared memory, the loads are fast.
Given that total memory is ~500KB and it's all in L2 cache after first access, and the kernel is tiny, the performance will be excellent.
Let me write the full Python module now.
```python
import torch
from torch.utils.cpp_extension import load_inline
cuda_source = """
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#define MAX_HEAD_DIM 128
__global__ void gated_delta_attention_kernel(
const float* __restrict__ q,
const float* __restrict__ k,
const float* __restrict__ v,
const float* __restrict__ gate,
float* __restrict__ out,
int B,
int T,
int H,
int D
) {
const int bh = blockIdx.x;
const int b = bh / H;
const int h = bh % H;
const int j = threadIdx.x;
float state[MAX_HEAD_DIM];
#pragma unroll
for (int i = 0; i < MAX_HEAD_DIM; i++) {
state[i] = 0.0f;
}
extern __shared__ float smem[];
float* k_smem = smem;
float* q_smem = smem + D;
for (int t = 0; t < T; t++) {
const int base = ((b * T + t) * H + h) * D;
k_smem[j] = k[base + j];
q_smem[j] = q[base + j];
const float gate_j = gate[base + j];
const float v_j = v[base + j];
__syncthreads();
#pragma unroll 4
for (int i = 0; i < D; i++) {
state[i] = gate_j * state[i] + k_smem[i] * v_j;
}
float dot = 0.0f;
#pragma unroll 4
for (int i = 0; i < D; i++) {
dot += q_smem[i] * state[i];
}
out[base + j] = dot;
__syncthreads();
}
}
"""
Wait, I changed to dynamic shared memory. But then k_smem[j] is written by thread j, which is fine. But with dynamic shared memory, we need to pass the size at kernel launch: `<<<num_blocks, num_threads, 2*D*sizeof(float)>>>`.
Also, the `#pragma unroll 4` with a runtime bound D means the compiler will unroll by 4 but still have a runtime loop. That's fine.
But one issue: `for (int i = 0; i < MAX_HEAD_DIM; i++) state[i] = 0.0f;` - this initializes all 128 elements, but we only need D. For D=24, this wastes some compute but is negligible.
Actually, better: `for (int i = 0; i < D; i++) state[i] = 0.0f;`
And we don't need MAX_HEAD_DIM for the template. But `float state[MAX_HEAD_DIM]` still declares 128 floats. We could make it smaller. But 128 is safe.
Actually, maybe I should just use a fixed-size shared memory array to avoid passing dynamic shared memory size. `__shared__ float k_smem[128]; __shared__ float q_smem[128];` But then we need a compile-time max. The dynamic approach is cleaner because we avoid wasting shared memory.
But wait, shared memory is per block and we're only using 2*D*4 bytes. It doesn't matter whether we allocate statically or dynamically. Dynamic is fine.
However, there's a small issue: if D > MAX_HEAD_DIM, our state array overflows. But the inputs are fixed at D=24. To be general, we could make MAX_HEAD_DIM = 256 or dispatch to a kernel with a dynamic state array in shared memory. But for now, MAX_HEAD_DIM = 128 is plenty.
Actually, let me think about whether we can put state in shared memory too and have a thread-block-level approach.
Alternative kernel:
- Each block handles one (b, h)
- Shared memory: `float state[D * D];` (because state is D x D)
- For each t:
- Load k_t[D], q_t[D], gate_t[D], v_t[D] into shared memory
- Each thread handles one (i, j) pair: total D*D threads
- Thread (i, j) updates `state[i*D + j] = gate_j * state[i*D + j] + k_i * v_j`
- Then output y[t, j] = sum_i q_i * state[i*D + j] for each j
- This requires a reduction over i for each j
This has D*D threads per block = 576 threads. With 8 blocks, total threads = 4608. That's much better occupancy! And each block has enough threads to fill an SM.
But the reduction is the tricky part. After updating state, we need to compute for each j: `sum_i q_i * state[i,j]`. With 576 threads, we can organize them as a 2D grid (D, D) within the block.
Let me think about this more carefully.
Block configuration: `dim3 blockDim(D, D);` where D=24. Total threads = 576.
Grid: `dim3 grid(B, H);`
Thread (i, j) in block (b, h):
- For each t:
- Load k_t[i] into k_smem[i] (all threads with same i participate - actually thread (i, j) can load k_t[i] but there are D threads with the same i, so we'd have redundant loads unless we assign loading to one thread per row)
- Better: in the first warp (or use some threads), load the 1D arrays.
Actually, for a 2D block of (D, D):
- Thread (i, j) loads: gate_j, v_j, k_i, q_i
- But k_i and q_i are the same for all j with the same i. So threads (i, 0..D-1) all try to load k_i and q_i. This is broadcast but redundant.
To avoid redundant loads, we can have only thread (i, 0) load k_i and q_i. But that complicates code.
With a simpler approach: all threads load what they need. The redundant loads will hit L1 cache and be fast.
After loading into shared memory:
```cpp
k_smem[i] = k[base + i]; // redundant load by all threads with same i
q_smem[i] = q[base + i]; // same
gate_smem[j] = gate[base + j];
v_smem[j] = v[base + j];
__syncthreads();
```
Then thread (i, j):
```cpp
state[ij] = gate_smem[j] * state[ij] + k_smem[i] * v_smem[j];
```
Where state[ij] is stored in a shared memory or register. If we store the entire D*D state in shared memory, we need D*D * 4 bytes = 2304 bytes. Plus k_smem, q_smem, gate_smem, v_smem = 4*D*4 = 384 bytes. Total ~2.7KB. Well within limits.
After updating state:
```cpp
__syncthreads();
// Now compute output: y[t, j] = sum_i q_i * state[i,j]
// Each thread (i, j) computes partial = q_i * state[i,j]
// Then reduce over i for each j
float partial = q_smem[i] * state[ij];
// Put partial into a reduction shared memory
// ...
```
For the reduction, we can use shared memory:
```cpp
// In shared memory, after state and k,q,gate,v arrays, we can have reduction array
// But we need to store partial products for the reduction.
// With 2D block (D, D), for each fixed j, we have D partial products (one per i).
// We can have an additional shared memory array `reduction_smem[D][D]` or just reuse.
```
Actually, we can do the reduction in place using warp shuffles if we organize carefully. But with 576 threads, we'd have 18 warps of 32 threads. The block is 24x24. Warps are filled in x-major order. So warp 0 contains threads (0..31). Since D=24, this includes threads (0,0..23), (1,0..7). So consecutive threads in a warp have varying j. This makes warp-level reduction tricky because we want to reduce over i for fixed j.
For fixed j, threads with the same j and different i are spaced D apart in thread ID: thread_id = i * D + j. So for j=0, threads are 0, 24, 48, ..., 528. These threads are not in the same warp (warp 0: 0-31; warp 1: 32-63, etc.).
- Thread 0 (i=0, j=0) is in warp 0
- Thread 24 (i=1, j=0) is in warp 0
- Thread 48 (i=2, j=0) is in warp 1
So the reduction would need cross-warp synchronization.
Maybe simpler: just have the 1D approach with 24 threads per block. It's cleaner, less synchronization, and the occupancy issue doesn't actually matter for a tiny problem.
Let me think about this differently. The PyTorch reference is clearly bottlenecked by Python overhead. A CUDA kernel with 8 blocks and 24 threads will likely run in under 20 microseconds. A Python loop with 128 iterations, each doing tensor operations, probably takes 500-2000 microseconds. So even with terrible GPU utilization, we'll get a huge speedup.
So let me just implement the clean 1D version and see.
Actually, I can increase parallelism by having each thread handle multiple j values. For example, with a block size of 128 threads and D=24, we can have each block handle 5 different (b, h) values or... wait, we only have 8 total (b, h).
What if we have each block handle one (b, h) and each of the 24 threads handles one j, but we launch with block size = 32 instead of 24? Then we have 8 threads doing nothing. Each block is one warp. The 8 inactive threads won't cause issues as long as they don't execute shared memory operations and `__syncthreads()`...
Wait, if block size is 32 and we only use threads 0-23:
- Threads 24-31 would execute `k_smem[j] = k[base + j]` with j=24..31 which is out of bounds! BAD.
- We'd need to guard with `if (j < D)`. But then threads 24-31 skip the writes and `__syncthreads()`. In CUDA, all threads in a block must execute `__syncthreads()`. If some threads skip it due to a conditional, it's undefined behavior and can deadlock.
However, since block size is 32, and the condition `j < D` means threads 24-31 don't reach the first `__syncthreads()`. This will deadlock. So block size MUST be exactly D, or D must equal block size.
If we want a block size of 32, we'd need to pad D to 32. But inputs are D=24.
So I must use block size = D. For D=24, that's fine.
Alternatively, I can remove `__syncthreads()` altogether! If I don't use shared memory and each thread loads its own k and q values directly, the loads will be cached but there's no need for thread synchronization. However, each thread would redundantly load k[i] and q[i] for all i, which is bad for bandwidth.
But wait, if we do this without shared memory and with just k and q loaded from global memory into registers:
```cpp
for (int i = 0; i < D; i++) {
float k_i = k[base + i]; // All threads load the same k_i (D redundant loads)
state[i] = gate_j * state[i] + k_i * v_j;
}
float dot = 0.0f;
for (int i = 0; i < D; i++) {
float q_i = q[base + i];
dot += q_i * state[i];
}
```
With 24 threads per block, each loading k[base + i] for i=0..23, we have 24*24 = 576 loads of floats for k, but only 24 unique values. This is 24x redundancy. Same for q. For a bandwidth-bound kernel, this could hurt performance significantly.
However, since the loads are all to the same address within a warp, they might be combined into a single load via L1 broadcast. In a single warp, when all threads read the same address, it's broadcast in one transaction. But here, thread j reads k[base + 0], k[base + 1], ..., k[base + 23] sequentially.
For thread j, in iteration i=0, all 24 threads read k[base + 0]. This is a broadcast - all threads in the warp read the same address, so it's one transaction. Then in iteration i=1, all threads read k[base + 1], again one transaction. So total transactions for k per step = D = 24. Same as if we loaded k once per block into shared memory. The difference is that instead of doing it parallel (coalesced) in one transaction, we do it sequentially in D transactions. But with a single warp and broadcast, it's similar.
Actually, if all threads in a warp simultaneously read k[base + i], it's one memory transaction. If we have a loop over i, we do one transaction per iteration. That's 24 transactions for the k loop. In the shared memory approach, we also do D loads (one per thread, coalesced), which might be done in one or two transactions. So shared memory saves a factor of ~D in transaction count.
For D=24, the shared memory approach does k in ~1-2 transactions, while the naive approach does it in 24 transactions. But 24 transactions for a 96-byte array is still very fast. In fact, for T=128, total k transactions per block = 128 * 24 = 3072 transactions. At 3.35 TB/s and each transaction ~32 bytes, that's ~0.03 ms. Still tiny.
So even the naive approach without shared memory would be fast vs the Python loop. But let's use shared memory to be safe.
Now, one important optimization: use `__restrict__` pointers.
Also, let me think about whether compiling with `-O3` is default for `load_inline`. Yes, it is.
Let me write the code and compile it.
One more thing: the `state` array. I used dynamic shared memory and had state in registers. Good.
For dynamic shared memory:
```cpp
extern __shared__ float smem[];
float* k_smem = smem;
float* q_smem = smem + D;
```
But D is a runtime parameter. Dynamic shared memory size is set at kernel launch: `<<<grid, block, 2 * D * sizeof(float)>>>`. That's fine.
Wait, but ` extern __shared__ float smem[];` in a kernel gets the size from the third kernel launch parameter. We set it to `2 * D * sizeof(float)`. Good.
Let me think about whether `out[base + j] = dot;` causes a RAW hazard or anything. No, each thread writes to a unique address.
Also, there's the issue that k_smem and q_smem are only D elements, and we do `k_smem[j] = ...` where j = threadIdx.x. Since block size = D, all threads 0..D-1 write to valid locations. Good.
One subtlety: if D is not equal to the warp size (32), then the warp has some threads that are valid and some that are "disabled" only in the sense that they don't exist because block size < 32. If block size is 24, then there are exactly 24 active threads in a warp. The other 8 lane positions are inactive. Is `__syncthreads()` still valid? Yes, because there are no threads in those lanes. All active threads (0..23) execute `__syncthreads()`. This is fine.
Wait, I just recalled: in CUDA, `__syncthreads()` requires all threads in the block to reach it. If block size is 24, all 24 active threads must reach it. There are no "missing" threads to worry about because the block is only 24 threads. A warp of 32 includes 24 active threads and 8 inactive ones, but inactive threads don't participate in any instructions. The `__syncthreads()` instruction is executed by the active threads in the warp. Since all 24 threads in the block reach it, it's fine. The inactive lanes don't matter because they don't exist at the block level.
So block size = D = 24 is perfectly fine.
Let me write the complete Python module.
Actually, I should also consider: `torch.zeros_like(q)` for the output. Then we call the kernel. The kernel computes all outputs, so zero initialization is unnecessary but harmless.
Here's the full module:
```python
import torch
from torch.utils.cpp_extension import load_inline
cuda_source = """
#include <torch/extension.h>
#include <cuda_runtime.h>
#define MAX_HEAD_DIM 128
__global__ void gated_delta_attention_kernel(
const float* __restrict__ q,
const float* __restrict__ k,
const float* __restrict__ v,
const float* __restrict__ gate,
float* __restrict__ out,
int B,
int T,
int H,
int D
) {
const int bh = blockIdx.x;
const int b = bh / H;
const int h = bh % H;
const int j = threadIdx.x;
float state[MAX_HEAD_DIM];
#pragma unroll
for (int i = 0; i < D; i++) {
state[i] = 0.0f;
}
extern __shared__ float smem[];
float* k_smem = smem;
float* q_smem = smem + D;
for (int t = 0; t < T; t++) {
const int base = ((b * T + t) * H + h) * D;
k_smem[j] = k[base + j];
q_smem[j] = q[base + j];
const float gate_j = gate[base + j];
const float v_j = v[base + j];
__syncthreads();
#pragma unroll 4
for (int i = 0; i < D; i++) {
state[i] = gate_j * state[i] + k_smem[i] * v_j;
}
float dot = 0.0f;
#pragma unroll 4
for (int i = 0; i < D; i++) {
dot += q_smem[i] * state[i];
}
out[base + j] = dot;
__syncthreads();
}
}
"""
cpp_source = """
torch::Tensor gated_delta_attention_forward(
torch::Tensor q,
torch::Tensor k,
torch::Tensor v,
torch::Tensor gate
) {
int B = q.size(0);
int T = q.size(1);
int H = q.size(2);
int D = q.size(3);
auto out = torch::empty_like(q);
int num_blocks = B * H;
int num_threads = D;
size_t smem_size = 2 * D * sizeof(float);
gated_delta_attention_kernel<<<num_blocks, num_threads, smem_size>>>(
q.data_ptr<float>(),
k.data_ptr<float>(),
v.data_ptr<float>(),
gate.data_ptr<float>(),
out.data_ptr<float>(),
B, T, H, D
);
return out;
}
"""
gated_delta = load_inline(
name="gated_delta_attention",
cpp_sources=cpp_source,
cuda_sources=cuda_source,
functions=["gated_delta_attention_forward"],
extra_cuda_cflags=["-O3"],
)
class ModelNew(torch.nn.Module):
def __init__(self):
super().__init__()
def forward(self, q, k, v, gate):
return gated_delta.gated_delta_attention_forward(q, k, v, gate)
```
Wait, I need to make sure `load_inline` is called at module level so it compiles once. And the module-level variable `gated_delta` is accessible inside `ModelNew`.
One issue: `#pragma unroll` with a runtime loop bound doesn't actually unroll the loop. The compiler can only unroll when the bound is constant. But `#pragma unroll 4` with a runtime bound means the compiler unrolls the body 4 times and keeps a runtime loop around it with iterations in steps of 4. But if D is not a compile-time constant, the compiler can't fully unroll.
For D=24 (a compile-time constant in the problem context), it would be nice to fully unroll. We can add a template-based dispatch to get full unrolling for small D values, especially D=24.
Let me add template dispatching. This might improve performance slightly.
```cpp
template <int D>
__global__ void gated_delta_attention_kernel(
const float* __restrict__ q,
const float* __restrict__ k,
const float* __restrict__ v,
const float* __restrict__ gate,
float* __restrict__ out,
int B,
int T,
int H
) {
const int bh = blockIdx.x;
const int b = bh / H;
const int h = bh % H;
const int j = threadIdx.x;
float state[D];
#pragma unroll
for (int i = 0; i < D; i++) {
state[i] = 0.0f;
}
extern __shared__ float smem[];
float* k_smem = smem;
float* q_smem = smem + D;
for (int t = 0; t < T; t++) {
const int base = ((b * T + t) * H + h) * D;
k_smem[j] = k[base + j];
q_smem[j] = q[base + j];
const float gate_j = gate[base + j];
const float v_j = v[base + j];
__syncthreads();
#pragma unroll
for (int i = 0; i < D; i++) {
state[i] = gate_j * state[i] + k_smem[i] * v_j;
}
float dot = 0.0f;
#pragma unroll
for (int i = 0; i < D; i++) {
dot += q_smem[i] * state[i];
}
out[base + j] = dot;
__syncthreads();
}
}
// General fallback
__global__ void gated_delta_attention_kernel_general(
const float* __restrict__ q,
const float* __restrict__ k,
const float* __restrict__ v,
const float* __restrict__ gate,
float* __restrict__ out,
int B,
int T,
int H,
int D
) {
// same as above but with MAX_HEAD_DIM
...
}
```
Then in the launcher:
```cpp
switch (D) {
case 8: gated_delta_attention_kernel<8><<<...>>>(...); break;
case 16: gated_delta_attention_kernel<16><<<...>>>(...); break;
case 24: gated_delta_attention_kernel<24><<<...>>>(...); break;
case 32: gated_delta_attention_kernel<32><<<...>>>(...); break;
case 48: gated_delta_attention_kernel<48><<<...>>>(...); break;
case 64: gated_delta_attention_kernel<64><<<...>>>(...); break;
case 96: gated_delta_attention_kernel<96><<<...>>>(...); break;
case 128: gated_delta_attention_kernel<128><<<...>>>(...); break;
default: gated_delta_attention_kernel_general<<<...>>>(...); break;
}
```
This gives full unrolling for common head dimensions. For D=24, the inner loops over i will be fully unrolled, which should improve register usage and instruction scheduling.
Let me write the full template version.
Actually, with dynamic shared memory and a template parameter D, we need to pass the smem size at launch. But D is a template parameter at launch time, so we can compute `2 * D * sizeof(float)` as a compile-time constant. Perfect.
Let me think about whether `float state[D]` with template parameter D is valid in CUDA. Yes, it's fine - the array size is known at compile time for a specific instantiation.
The size of shared memory with template parameter: `extern __shared__ float smem[];` ... this is still dynamic shared memory. But we know the size at kernel launch: we pass `2 * D * sizeof(float)`. Good.
One more thing: in the template version, k_smem is `smem`, q_smem is `smem + D`. Thread j writes to k_smem[j] and q_smem[j]. Since j < D (block size = D), this is always valid.
Now, is there any issue with `out` not being contiguous in a way that matters? The output is (B, T, H, D), same as inputs. So `base + j` is the correct index.
OK, I think this is good. Let me compile and test it.
Wait, I should make sure the load_inline call is correct. I need:
```python
gated_delta_op = load_inline(
name="gated_delta_attention",
cpp_sources=cpp_source,
cuda_sources=cuda_source,
functions=["gated_delta_attention_forward"],
extra_cuda_cflags=["-O3", "--use_fast_math"],
)
```
`--use_fast_math` might change numerical results slightly due to approximate division and sqrt, but we only do multiplication and addition, so it shouldn't matter. I can add it for extra optimization, but without it is safer for correctness. I'll try without first.
Also, I should check: does `load_inline` handle multiple CUDA source strings correctly? Yes, they are concatenated.
Let me write the complete module now.
One more thought: since the kernel is very small, maybe we should use a fused kernel or consider persistent threads. But again, overkill.
Actually, I realize that q_smem and k_smem are loaded at every step. We could potentially prefetch the next step while computing the current step using more shared memory. But with T=128 and tiny arrays, L1 cache probably handles this.
Let's just write and compile.
Another optimization: `__ldg` for the loads from k, q, gate, v. But it's fine without.
Oh, one more thing! `q`, `k`, `v`, `gate` need to be contiguous in the last dimension for our indexing to work. PyTorch tensors by default have row-major (last dimension contiguous) for float32. But if someone passes non-contiguous tensors, we'd have issues. In the reference implementation, PyTorch operations handle strided tensors. Our kernel assumes contiguous memory layout.
To be safe, we should check if the tensors are contiguous and call `.contiguous()` if not. Or we can just pass `.contiguous()` always, but that might add overhead. Better to call `.contiguous()` on inputs before passing to the kernel.
But the test probably uses `torch.randn` which gives contiguous tensors. Still, it's good practice to handle it:
```python
q = q.contiguous()
k = k.contiguous()
v = v.contiguous()
gate = gate.contiguous()
```
But `.contiguous()` might copy if non-contiguous. For `forward()` in ModelNew, we can just call `.contiguous()` on all inputs. The test probably uses contiguous tensors, so it will be a no-op.
Actually, looking at the reference, there are no explicit contiguous calls. But `torch.randn` creates contiguous tensors. `torch.sigmoid` preserves contiguity. So we're fine.
Let me also implement a CPU fallback path? The kernel is CUDA-only. The test likely only runs on CUDA. The problem says "write a custom CUDA kernel" and "faster than the PyTorch reference", implying GPU environment.
But to be safe, if inputs are on CPU, we should fall back to PyTorch. However, the inputs are created with `torch.randn(...)` without device specification, so they might be on CPU by default... wait, the reference model is a nn.Module and the test probably moves tensors to CUDA or runs them on CPU. But the task specifically says "custom CUDA kernel", so GPU is available.
Let me just assume GPU.
One more detail: the kernel should be robust to any batch size, seq len, etc. But we can assume standard sizes.
Let me write the final code now.
```python
import torch
from torch.utils.cpp_extension import load_inline
cuda_source = """
#include <torch/extension.h>
#include <cuda_runtime.h>
template <int D>
__global__ void gated_delta_attention_kernel(
const float* __restrict__ q,
const float* __restrict__ k,
const float* __restrict__ v,
const float* __restrict__ gate,
float* __restrict__ out,
int B,
int T,
int H
) {
const int bh = blockIdx.x;
const int b = bh / H;
const int h = bh % H;
const int j = threadIdx.x;
float state[D];
#pragma unroll
for (int i = 0; i < D; ++i) {
state[i] = 0.0f;
}
extern __shared__ float smem[];
float* k_smem = smem;
float* q_smem = smem + D;
for (int t = 0; t < T; ++t) {
const int base = ((b * T + t) * H + h) * D;
k_smem[j] = k[base + j];
q_smem[j] = q[base + j];
const float gate_j = gate[base + j];
const float v_j = v[base + j];
__syncthreads();
#pragma unroll
for (int i = 0; i < D; ++i) {
state[i] = gate_j * state[i] + k_smem[i] * v_j;
}
float dot = 0.0f;
#pragma unroll
for (int i = 0; i < D; ++i) {
dot += q_smem[i] * state[i];
}
out[base + j] = dot;
__syncthreads();
}
}
#define DISPATCH_D(D, ...) \
[&] { \
switch (D) { \
case 8: gated_delta_attention_kernel<8> <<<__VA_ARGS__>>>; break; \
case 16: gated_delta_attention_kernel<16><<<__VA_ARGS__>>>; break; \
case 24: gated_delta_attention_kernel<24><<<__VA_ARGS__>>>; break; \
case 32: gated_delta_attention_kernel<32><<<__VA_ARGS__>>>; break; \
case 48: gated_delta_attention_kernel<48><<<__VA_ARGS__>>>; break; \
case 64: gated_delta_attention_kernel<64><<<__VA_ARGS__>>>; break; \
case 96: gated_delta_attention_kernel<96><<<__VA_ARGS__>>>; break; \
case 128: gated_delta_attention_kernel<128><<<__VA_ARGS__>>>; break; \
default: \
TORCH_CHECK(false, "Unsupported head_dim: ", D); \
} \
}()
"""
cpp_source = """
torch::Tensor gated_delta_attention_forward(
torch::Tensor q,
torch::Tensor k,
torch::Tensor v,
torch::Tensor gate
) {
int B = q.size(0);
int T = q.size(1);
int H = q.size(2);
int D = q.size(3);
TORCH_CHECK(q.is_cuda(), "q must be a CUDA tensor");
TORCH_CHECK(k.is_cuda(), "k must be a CUDA tensor");
TORCH_CHECK(v.is_cuda(), "v must be a CUDA tensor");
TORCH_CHECK(gate.is_cuda(), "gate must be a CUDA tensor");
auto out = torch::empty_like(q);
int num_blocks = B * H;
int num_threads = D;
size_t smem_size = 2 * D * sizeof(float);
DISPATCH_D(D,
num_blocks, num_threads, smem_size,
q.data_ptr<float>(),
k.data_ptr<float>(),
v.data_ptr<float>(),
gate.data_ptr<float>(),
out.data_ptr<float>(),
B, T, H
);
return out;
}
"""
```
Wait, the `DISPATCH_D` macro with `__VA_ARGS__` is a bit tricky with the kernel launch syntax `<<<...>>>`. The macro expansion needs to be `kernel<<<grid, block, smem>>>(args)`. My macro usage:
```cpp
DISPATCH_D(D,
num_blocks, num_threads, smem_size,
q.data_ptr<float>(),
...
);
```
This would expand to:
```cpp
gated_delta_attention_kernel<24><<<num_blocks, num_threads, smem_size>>>(q.data_ptr<float>(), ...);
```
But I need to make sure the macro expansion works correctly. Let me use a simpler macro:
```cpp
#define DISPATCH_D(D, num_blocks, num_threads, smem_size, ...) \
do { \
if (D == 8) { \
gated_delta_attention_kernel<8><<<num_blocks, num_threads, smem_size>>>(__VA_ARGS__); \
} else if (D == 16) { \
gated_delta_attention_kernel<16><<<num_blocks, num_threads, smem_size>>>(__VA_ARGS__); \
} else if (D == 24) { \
gated_delta_attention_kernel<24><<<num_blocks, num_threads, smem_size>>>(__VA_ARGS__); \
} else if (D == 32) { \
gated_delta_attention_kernel<32><<<num_blocks, num_threads, smem_size>>>(__VA_ARGS__); \
} else if (D == 48) { \
gated_delta_attention_kernel<48><<<num_blocks, num_threads, smem_size>>>(__VA_ARGS__); \
} else if (D == 64) { \
gated_delta_attention_kernel<64><<<num_blocks, num_threads, smem_size>>>(__VA_ARGS__); \
} else if (D == 96) { \
gated_delta_attention_kernel<96><<<num_blocks, num_threads, smem_size>>>(__VA_ARGS__); \
} else if (D == 128) { \
gated_delta_attention_kernel<128><<<num_blocks, num_threads, smem_size>>>(__VA_ARGS__); \
} else { \
TORCH_CHECK(false, "Unsupported head_dim: ", D); \
} \
} while(0)
```
This is straightforward and works. The macro is placed in the cuda_source so it's available in the compilation unit that contains the kernel.
Wait, `load_inline` puts `cuda_sources` and `cpp_sources` in separate compilation units by default? Let me check. Actually, `load_inline` concatenates all CUDA sources into a .cu file and all C++ sources into a .cpp file. But with `cuda_sources` and `cpp_sources`, it compiles them separately and links them.
If the kernel declaration is in the .cu file and the dispatcher is in the .cpp file, the .cpp file won't see the kernel declaration! We need the dispatcher to be in the .cu file or have a declaration.
`load_inline` puts all CUDA code in one `.cu` and all C++ in one `.cpp`. For the C++ dispatcher to call the CUDA kernel, it needs the kernel to be declared in a header or the C++ file needs to see it. But `.cu` files are compiled by nvcc and `.cpp` by the regular C++ compiler. The kernel function is compiled to device code and callable from host code in the same compilation unit.
Actually, CUDA kernels can be declared in `.cu` and called from host code in the same `.cu` file. But if the host caller is in `.cpp`, it typically needs to be compiled by nvcc as well, or there needs to be a wrapper.
`load_inline` creates:
- `.cu` containing all `cuda_sources`
- `.cpp` containing all `cpp_sources`
- It compiles both and links them.
But calling a `__global__` function from `.cpp` compiled by g++ might not work because the kernel launch syntax `<<<>>>` is not valid C++. Actually, in modern PyTorch extension compilation, `.cpp` files are compiled as C++ and can't contain CUDA kernel launch syntax. The `cpp_sources` are compiled with `c++` not `nvcc` for `.cpp` files.
So I should put the dispatcher in the `cuda_sources` as a regular C++ function, and then call that function from `cpp_sources`. Or put everything in `cuda_sources`.
Actually, looking at the docs for `load_inline`:
- The code is split into C++ and CUDA parts
- `functions` should list the C++ host functions that are exposed
The typical pattern is:
- Define `__global__` kernels in `cuda_sources`
- Define host-side wrapper functions in `cpp_sources` (or `cuda_sources`) that call the kernels
- List the wrapper functions in `functions`
But the host-side wrapper can be in `cuda_sources` because nvcc compiles it and can handle kernel launches. Or it can be in `cpp_sources` if the kernel is declared with `extern "C"` or just declared without `__global__` attribute? Actually, calling a `__global__` from a `.cpp` file is tricky.
Let me just put everything in `cuda_sources` and have the wrapper function there too. Then `cpp_sources` can be empty or minimal. Since `load_inline` requires `cpp_sources` to be a string, we can put the wrapper there. Wait, but kernel calls in `.cpp` with `<<<>>>` won't compile.
The safest approach is to put both the kernel and the wrapper in `cuda_sources`. The `cpp_sources` can be empty. And we list the wrapper function in `functions`. But the wrapper is in the `.cu` file, which is compiled by nvcc.
Actually, looking at examples from PyTorch docs:
```python
load_inline(
name="my_op",
cpp_sources="torch::Tensor my_op(torch::Tensor x) { ... }",
cuda_sources="__global__ void kernel(...) { ... }",
functions=["my_op"]
)
```
In this pattern, `my_op` calls `kernel<<<>>>`. But if `my_op` is in `cpp_sources`, how does it call the kernel? The answer is that `load_inline` compiles all sources together in a way that they can reference each other. Actually, looking at the source code of `load_inline`, it concatenates `cpp_sources` and `cuda_sources` and compiles with nvcc.
Wait, let me think. In `torch.utils.cpp_extension.load_inline`, for the `sources` parameter, it handles `.cpp` and `.cu` internally. But when `cpp_sources` and `cuda_sources` are separate, it might compile them as separate files.
If `my_op` in `cpp_sources` calls a kernel, the kernel needs forward declaration. But the `<<<>>>` syntax is only valid in CUDA compilation units. So if `cpp_sources` is compiled as C++, it will fail.
I think the standard approach is: put everything in `cuda_sources` and leave `cpp_sources` empty, or put everything in `cpp_sources` where C++ contains a function declaration and the actual definition with launch code is in a `.cu` file. But for `load_inline`, it's simpler to just put the host wrapper that does `kernel<<<...>>>` in the same source as the kernel.
Actually, many examples online do this:
```python
source = """
#include <torch/extension.h>
__global__ void kernel(...) { ... }
torch::Tensor forward(...) {
kernel<<<...>>>(...);
return out;
}
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
m.def("forward", &forward, "...");
}
"""
```
But in this task, it says: "Do not write your own `PYBIND11_MODULE(...)` block in `cpp_sources`. `load_inline` auto-generates one from the `functions=[...]` argument"
So the pattern should be:
- `cpp_sources` contains only the C++ wrapper function (with kernel forward declaration if needed)
- `cuda_sources` contains the CUDA kernel
But can `cpp_sources` call a kernel with `<<<...>>>`? Only if it's compiled by nvcc. PyTorch's `load_inline` might compile `cpp_sources` with nvcc if CUDA is available?
Let me check torch.utils.cpp_extension behavior. According to PyTorch docs and source:
- `load_inline` compiles all C++ and CUDA code together.
- It manages compilation such that functions are visible across compilation units.
Actually, the cleanest way is to put the wrapper function in `cpp_sources` with a forward declaration of the kernel:
```cpp
// In cpp_sources
void my_kernel(...); // forward decl
torch::Tensor my_func(...) {
my_kernel<<<grid, block>>>(...);
return out;
}
```
And define the kernel in `cuda_sources`:
```cuda
__global__ void my_kernel(...) { ... }
```
This works because nvcc compiles the `.cpp` file as CUDA-C++ (with `-x cu` perhaps). Actually no, `.cpp` is compiled with the C++ compiler, not nvcc. So `<<<...>>>` in `.cpp` would fail.
Let me just put the host wrapper in the `cuda_sources` string. Then list it in `functions`. And make `cpp_sources` empty or minimal.
Wait, but `functions` lists functions from `cpp_sources`? Or from all sources? Actually, the auto-generated pybind11 wrapper binds the functions listed in `functions`. These functions could be in any source file. But typically the C++ functions are in `cpp_sources`.
Let me look at a reliable pattern:
```python
source = """
#include <torch/extension.h>
__global__ void kernel(...) { ... }
torch::Tensor my_forward(torch::Tensor x) {
...
kernel<<<grid, block>>>(...);
return y;
}
"""
load_inline(
name="my_ext",
cpp_sources="", # or source,
cuda_sources=source,
functions=["my_forward"],
extra_cuda_cflags=["-O3"]
)
```
Putting the host wrapper in `cuda_sources` should work because nvcc can compile both the `__global__` kernel and the regular C++ function that launches it.
Actually, `load_inline` might not let you specify the same content in both. Let me just put the kernel and the wrapper in `cuda_sources`, and leave `cpp_sources` empty.
Wait, looking at the signature:
```python
load_inline(name, cpp_sources, cuda_sources, functions, ...)
```
Both `cpp_sources` and `cuda_sources` are strings. The resulting extension will compile the cpp code and cuda code and link them. If I put the wrapper in `cuda_sources`, it should be visible to pybind11 because `functions` refers to it by name. But `functions` might expect the functions to be in the C++ binding layer?
Actually no, pybind11 binds any function with the given name that has external linkage. It doesn't matter if it's defined in `.cpp` or `.cu`.
Let me test this approach. I'll put the kernel in `cuda_sources` along with the wrapper function.
Actually, to be extra safe, I'll put the wrapper in `cpp_sources` with a forward declaration of the kernel, and put the kernel definition in `cuda_sources`. But then `cpp_sources` would have kernel launch syntax. Unless... I put the wrapper in `cuda_sources` too.
Let me try: all code in `cuda_sources`. Empty `cpp_sources`.
Wait, what if `cpp_sources` is empty string? That's probably fine.
Let me just create a comprehensive `cuda_sources` string containing the kernel and wrapper, and set `cpp_sources=""`.
But I want to be careful: sometimes empty strings cause issues. Let me set `cpp_sources=" "` or just include a minimal comment.
Let me look up `load_inline` behavior. In PyTorch's `cpp_extension.py`, `load_inline` internally does:
```python
cpp_sources = [cpp_sources] if isinstance(cpp_sources, str) else list(cpp_sources)
cuda_sources = [cuda_sources] if isinstance(cuda_sources, str) else list(cuda_sources)
# ... builds them into a tmpdir and calls load
```
So empty string is fine - it becomes an empty `.cpp` or `.cu` file.
I'll put wrapper in `cuda_sources`. Let's proceed.
Wait, actually for cleanliness, I can put the wrapper in `cpp_sources` using a helper that calls into CUDA code. But to avoid complexity, I'll just put everything in one file via `cuda_sources`.
Actually let me reconsider. `cpp_sources=" "` with just a space, and `cuda_sources` has everything.
But wait - in my macro approach, if I use `do { ... } while(0)` with `TORCH_CHECK` in a `.cu` file, that's fine because nvcc compiles it.
Let me write it cleanly.
Actually, maybe I should just avoid the template dispatch and use a fixed-size state array. The performance difference between unrolled and non-unrolled loops for D=24 is probably negligible compared to Python overhead and kernel launch cost. It also makes the code simpler and compilation faster.
For simplicity and robustness:
- Use a non-template kernel with `float state[128]` or similar
- Let the compiler handle the loop
- No dispatch macro needed
This reduces compilation time and complexity. Given that the performance target is "faster than PyTorch reference" (which is very slow due to Python loop), this simpler kernel will easily pass.
But I want to maximize speedup. Let me keep the template for D=24 since that's the test size, but I'll also add a fallback for completeness.
Actually, you know what, let me just make the code as simple and correct as possible first, verify it works, and then optimize.
Final code without dispatch:
```python
import torch
from torch.utils.cpp_extension import load_inline
cuda_source = """
#include <torch/extension.h>
#include <cuda_runtime.h>
#define MAX_HEAD_DIM 256
__global__ void gated_delta_attention_kernel(
const float* __restrict__ q,
const float* __restrict__ k,
const float* __restrict__ v,
const float* __restrict__ gate,
float* __restrict__ out,
int B,
int T,
int H,
int D
) {
int bh = blockIdx.x;
int b = bh / H;
int h = bh % H;
int j = threadIdx.x;
float state[MAX_HEAD_DIM];
for (int i = 0; i < D; i++) {
state[i] = 0.0f;
}
extern __shared__ float smem[];
float* k_smem = smem;
float* q_smem = smem + D;
for (int t = 0; t < T; t++) {
int base = ((b * T + t) * H + h) * D;
k_smem[j] = k[base + j];
q_smem[j] = q[base + j];
float gate_j = gate[base + j];
float v_j = v[base + j];
__syncthreads();
for (int i = 0; i < D; i++) {
state[i] = gate_j * state[i] + k_smem[i] * v_j;
}
float dot = 0.0f;
for (int i = 0; i < D; i++) {
dot += q_smem[i] * state[i];
}
out[base + j] = dot;
__syncthreads();
}
}
torch::Tensor gated_delta_attention_forward(
torch::Tensor q,
torch::Tensor k,
torch::Tensor v,
torch::Tensor gate
) {
int B = q.size(0);
int T = q.size(1);
int H = q.size(2);
int D = q.size(3);
auto out = torch::empty_like(q);
int num_blocks = B * H;
int num_threads = D;
size_t smem_size = 2 * D * sizeof(float);
gated_delta_attention_kernel<<<num_blocks, num_threads, smem_size>>>(
q.data_ptr<float>(),
k.data_ptr<float>(),
v.data_ptr<float>(),
gate.data_ptr<float>(),
out.data_ptr<float>(),
B, T, H, D
);
return out;
}
"""
gated_delta_attention = load_inline(
name="gated_delta_attention",
cpp_sources="",
cuda_sources=cuda_source,
functions=["gated_delta_attention_forward"],
extra_cuda_cflags=["-O3"],
)
class ModelNew(torch.nn.Module):
def __init__(self):
super().__init__()
def forward(self, q, k, v, gate):
return gated_delta_attention.gated_delta_attention_forward(q, k, v, gate)
```
Wait, if `cpp_sources=""`, and `functions=["gated_delta_attention_forward"]` refers to a function in `cuda_sources`, will pybind11 find it? Let me think...
The auto-generated bindings likely do something like:
```cpp
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
m.def("gated_delta_attention_forward", &::gated_delta_attention_forward, ...);
}
```
This requires `gated_delta_attention_forward` to have external linkage and be visible. Since it's defined in the `.cu` file, it should be visible to the linker as long as it's not static. Since it's defined in `cuda_sources` and the pybind module is auto-generated, the linker should find it. This should work.
Actually, I recall that `load_inline` might only scan `cpp_sources` for `functions`. But according to the docs: "functions (List[str]) – a list of function names for which to generate bindings." It doesn't specify they must be in `cpp_sources`. They can be in any source.
I'll try this. If it doesn't work, I'll move the wrapper to `cpp_sources` and keep only the `__global__` kernel in `cuda_sources`.
But if I move the wrapper to `cpp_sources`, the kernel launch syntax `<<<>>>` won't compile with a regular C++ compiler. Unless the wrapper is compiled by nvcc. Some build systems compile `.cpp` with nvcc when building CUDA extensions. But `load_inline` specifically separates them.
Looking at PyTorch source code for `load_inline`: in `_jit_compile`, it writes the sources to files with their respective extensions and calls `load`. The `load` function creates a `BuildExtension` that uses `nvcc` for `.cu` and the regular C++ compiler for `.cpp`. So `<<<...>>>` in `.cpp` would indeed fail.
So the wrapper MUST be in the `.cu` file (or the `.cpp` must be compiled as CUDA, which it isn't).
Therefore, putting the wrapper in `cuda_sources` is correct.
Let me compile and test this now.
Actually, let me double-check that `state[MAX_HEAD_DIM]` with MAX_HEAD_DIM=256 doesn't cause register spilling issues. Each thread uses 256 floats = 1024 bytes of local memory if it spills. But we only access indices 0..D-1 where D=24. A smart compiler might allocate only 24 registers. But with `float state[256]`, the compiler sees an array of 256 and might allocate them all on the stack (local memory). This would be bad!
To avoid this, we can either:
1. Use a smaller MAX_HEAD_DIM like 64 or 32
2. Use template parameter D (fully unrolled, compiler knows only D registers needed)
3. Use shared memory for state
For D=24, if we declare `float state[64]`, the compiler might still put it in local memory. But 64 floats = 256 bytes. For 24 threads per block, total local memory per block = 24 * 256 = 6KB. Still fine.
But local memory is L1 cached and as fast as global memory. For a kernel doing 24 iterations, reading/writing local memory 24 times might add some latency. However, if the compiler is smart, it can keep the active elements in registers even with a declared array size of 64.
Well, there is a risk. To be safe, let me use template parameter D. Or shared memory.
Actually, let me use shared memory for state:
```cpp
extern __shared__ float smem[];
float* k_smem = smem;
float* q_smem = smem + D;
float* state_smem = smem + 2 * D + threadIdx.x * D;
```
Wait, no. If state is in shared memory, all threads can see each other's state. We only want each thread to have its own state.
We can structure shared memory as:
```cpp
extern __shared__ float smem[];
float* k_smem = smem;
float* q_smem = smem + D;
float* state = smem + 2 * D + threadIdx.x * D; // D floats per thread
```
Then `state[i]` accesses `smem[2*D + threadIdx.x*D + i]`. Since threadIdx.x ranges from 0 to D-1, and each thread accesses its own row, there are NO bank conflicts as long as each thread accesses its own contiguous chunk (which it does).
Shared memory size = 2*D + D*D floats = D*(D+2)*4 bytes. For D=24: 24*26*4 = 2496 bytes. Very manageable.
This approach avoids the register/local memory issue entirely and is very clean.
Kernel with shared memory for state:
```cpp
extern __shared__ float smem[];
float* k_smem = smem;
float* q_smem = smem + D;
float* my_state = smem + 2 * D + j * D;
for (int i = 0; i < D; i++) {
my_state[i] = 0.0f;
}
for (int t = 0; t < T; t++) {
int base = ((b * T + t) * H + h) * D;
k_smem[j] = k[base + j];
q_smem[j] = q[base + j];
float gate_j = gate[base + j];
float v_j = v[base + j];
__syncthreads();
for (int i = 0; i < D; i++) {
my_state[i] = gate_j * my_state[i] + k_smem[i] * v_j;
}
float dot = 0.0f;
for (int i = 0; i < D; i++) {
dot += q_smem[i] * my_state[i];
}
out[base + j] = dot;
__syncthreads();
}
```
Launch with smem_size = (2 * D + D * D) * sizeof(float).
Wait, but all threads write to `my_state[i]` and read only their own `my_state`. There are no bank conflicts.
Actually, threads write to `my_state[i] = ...` which is at address `smem[2*D + j*D + i]`. Since j is the thread ID and each thread accesses a different row, and i is sequential within the thread, there are no cross-thread accesses to the same bank (unless D causes wrapping, but with 32 banks and thread j accessing offset j*D, bank = (j*D + i) % 32. For D=24: thread 0 accesses banks 0,1,2,...,23. Thread 1 accesses banks 24,25,26,...,47%32=15, so banks wrap around. But since each thread accesses different memory locations and threads don't share state accesses, bank conflicts only occur if two threads access the same bank simultaneously.
Thread j accesses `j*D + i`. Thread j+1 accesses `(j+1)*D + i'` simultaneously (since they're in the same warp). Are these in the same bank? We need `(j*D + i) % 32 == ((j+1)*D + i') % 32`.
For the state initialization loop `for (int i = 0; i < D; i++)`, all threads have i as the loop variable. So for a given i, thread j writes to bank `(j*D + i) % 32`. Thread j+1 writes to bank `((j+1)*D + i) % 32`.
For D=24: thread j writes to bank `(24j + i) % 32 = (-8j + i) % 32`. Thread j+1 writes to bank `(-8(j+1) + i) % 32 = (-8j - 8 + i) % 32`. These are 8 banks apart (mod 32). Since they're different banks, no conflict. In fact, the stride is D=24. gcd(24, 32) = 8. So accesses in this pattern cycle through 32/8 = 4 banks? Wait...
Actually, if stride is 24 and we have 32 banks, and bank index = (address / 4) % 32:
- Thread 0: address = 0*24 + i, bank = (6i) % 32
- Thread 1: address = 1*24 + i, bank = (6 + 6i) % 32
- Thread 2: address = 2*24 + i, bank = (12 + 6i) % 32
- Thread 3: address = 3*24 + i, bank = (18 + 6i) % 32
- Thread 4: address = 4*24 + i, bank = (24 + 6i) % 32 = (-8 + 6i) % 32
- Thread 5: address = 5*24 + i, bank = (30 + 6i) % 32
- Thread 6: address = 6*24 + i, bank = (36 + 6i) % 32 = (4 + 6i) % 32
- Thread 7: address = 7*24 + i, bank = (42 + 6i) % 32 = (10 + 6i) % 32
For a fixed i, threads 0-7 access banks: 6i, 6+6i, 12+6i, 18+6i, -8+6i, 30+6i, 4+6i, 10+6i. All distinct! And there are 24 threads = 3 sets of 8 threads. Each set accesses distinct banks. So 0 bank conflicts!
Wait, stride is 24 floats = 96 bytes. Bank = (address_in_bytes / 4) % 32 = (96/4)j % 32 + (i) % 32 if i is in floats... wait, address in bytes is (2*D + j*D + i) * 4. Let me compute bank more carefully.
Bank index = (byte_address / 4) % 32 = ((2*D + j*D + i)) % 32.
For D=24: bank = (48 + 24j + i) % 32 = (16 + 24j + i) % 32 = (16 - 8j + i) % 32.
For fixed i and consecutive threads j=0,1,2,3...:
- j=0: (16+i)%32
- j=1: (8+i)%32
- j=2: (0+i)%32
- j=3: (-8+i)%32 = (24+i)%32
- j=4: (16+i)%32
- j=5: (8+i)%32
- j=6: (0+i)%32
- j=7: (24+i)%32
Oh! For j=0 and j=4: same bank! Similarly j=1 and j=5. So there ARE bank conflicts.
In a warp of 32 threads (or our 24 threads), threads that are 4 apart have the same bank. But they access different i at different times. Let's check if they conflict in a single instruction (same i):
For i fixed, threads j and j+4 both want to write to bank (16 + i) % 32 if j multiple of 4? Wait:
j=0: (16+i)%32
j=4: (16-32+i)%32 = (16+i)%32
Yes, for the same i, threads j and j+4 have the same bank index. So 8 threads (j=0,4,8,12,16,20) all access bank (16+i)%32 simultaneously. That's a 6-way bank conflict! (actually threads 0,4,8,12,16,20 and possibly 24 doesn't exist since D=24)
So there ARE bank conflicts in shared memory.
But wait - are these accesses happening simultaneously? In the initialization loop `for (int i = 0; i < D; i++) { my_state[i] = 0.0f; }`, all threads in the warp execute the same instruction at the same time (same i). So yes, for each i, threads 0,4,8,12,16,20 all write to the same bank. This is bad.
However, in the actual computation loop:
```cpp
for (int i = 0; i < D; i++) {
my_state[i] = gate_j * my_state[i] + k_smem[i] * v_j;
}
```
Here too, all threads execute with the same i simultaneously (lockstep within a warp). So we'd have the same bank conflict issue.
Actually, wait. Do we have bank conflicts on my_state[i]? Each thread accesses its own my_state[i], which is at a different address. Bank conflicts happen when multiple threads in the same warp access different addresses that map to the SAME bank. Yes, that's exactly what's happening.
For D=24 and bank = (address_in_bytes / 4) % 32:
- Thread 0 accesses address offset: 0*24 + i = i, bank = i % 32
- Thread 1 accesses address offset: 1*24 + i = 24+i, bank = (24+i) % 32 = (-8+i) % 32
- Thread 2 accesses: 48+i, bank = (48+i)%32 = (16+i)%32
- Thread 3 accesses: 72+i, bank = (72+i)%32 = (8+i)%32
- Thread 4 accesses: 96+i, bank = (96+i)%32 = i%32
Yes, thread 0 and thread 4 have the same bank for the same i. Since there are 24 threads, the pattern repeats every 4 threads: 6 groups of 4 threads each accessing the same bank.
Wait, no. Let me list all 24 threads for i=0:
j=0: bank 0
j=1: bank 24
j=2: bank 16
j=3: bank 8
j=4: bank 0
j=5: bank 24
j=6: bank 16
j=7: bank 8
j=8: bank 0
j=9: bank 24
j=10: bank 16
j=11: bank 8
j=12: bank 0
j=13: bank 24
j=14: bank 16
j=15: bank 8
j=16: bank 0
j=17: bank 24
j=18: bank 16
j=19: bank 8
j=20: bank 0
j=21: bank 24
j=22: bank 16
j=23: bank 8
So there are only 4 unique banks accessed: 0, 8, 16, 24. And 6 threads access each bank. This is a 6-way bank conflict. Ouch.
To avoid bank conflicts, we could pad the row size to avoid multiples that share banks. The bank width is 4 bytes, and there are 32 banks. Conflicts occur because D=24 and the stride is 96 bytes. The bank index increments by 96/4 = 24 mod 32 each row.
To avoid conflicts, we can use a padded row size that is coprime with 32, or at least doesn't create small periods. For example, pad to 25: stride = 100 bytes, bank stride = 25 mod 32. gcd(25, 32) = 1, so all threads access different banks.
Or pad to 33 (but that's wasteful). We can store state in shared memory with a padded row pitch.
```cpp
#define PITCH 33 // or whatever makes gcd(pitch, 32) = 1
extern __shared__ float smem[];
float* my_state = smem + 2 * D + j * PITCH; // each thread gets PITCH floats
```
For D=24 and PITCH=33: total shared memory for state = D * PITCH * 4 bytes = 24*33*4 = 3168 bytes. Plus 2*D*4 = 192 bytes. Total 3360 bytes. Still fine.
But then we need to make sure we don't waste too much shared memory. For D up to 128, 128*33*4 = 16896 bytes. Fine.
Wait, there's a much simpler option: keep state in registers! If each thread stores its own state vector in a register array, there are no shared memory bank conflict concerns. The issue was whether the compiler puts the array in local memory.
With `float state[24]` (or fixed max), a modern NVIDIA compiler (nvcc) for an array indexed by a compile-time constant or a small loop variable usually keeps it in registers. Since D=24 and each thread accesses `state[i]` where i is a loop variable 0..23, the compiler can map each `state[i]` to a separate register. So we'd have 24 registers per thread.
The catch was using `float state[MAX_HEAD_DIM]` with MAX_HEAD_DIM=256. The compiler might not be smart enough to realize only 24 are used. But `float state[24]` (or templated `float state[D]`) tells the compiler exactly.
Actually, here is the key: `float state[24]` where the indices used are `i` for `i < D` and D is a runtime parameter. The compiler sees the array size is 24 but doesn't know D. However, in the loop `for (int i = 0; i < D; i++)`, D is a runtime variable. The compiler likely puts the entire array in local memory because it can't prove the accesses are to fixed elements.
But wait, if we write:
```cpp
float s0, s1, s2, ..., s23;
```
Then manually loop, the compiler would put them in registers. But that's ugly.
Alternatively, with `#pragma unroll` and template parameter D, the compiler can see the fixed size and unroll the loop, assigning each `state[i]` to a register.
So template parameter D is much better for register allocation.
Let me switch to a template-based approach. Since D=24 is the test case, I'll instantiate for D=24. For other values, I'll fall back to shared memory or a generic kernel.
Wait, looking at the reference problem, `head_dim = 24`. But the Model doesn't have this as a fixed constant - it just reads from the input shape. The test might vary the shape, though. To be safe, I should handle arbitrary D.
For arbitrary D, I can write a general kernel that stores state in shared memory to avoid register pressure, and a specialized kernel for D=24 that stores state in registers.
Or, I can just write the template with D=24 as the primary path and a fallback using shared memory.
Actually, the simplest robust solution that also gets good performance is:
1. Template the kernel on D
2. Use `float state[D]` in registers
3. Add a dispatcher for common D values
4. For arbitrary D, write a version using a dynamically sized local array or shared memory
But for the fallback with arbitrary D > 64 or so, register pressure might be an issue. We can use shared memory for the fallback.
Let me do this:
```cpp
template <int D>
__global__ void gated_delta_attention_kernel_reg(
const float* __restrict__ q,
const float* __restrict__ k,
const float* __restrict__ v,
const float* __restrict__ gate,
float* __restrict__ out,
int B,
int T,
int H
) {
int bh = blockIdx.x;
int b = bh / H;
int h = bh % H;
int j = threadIdx.x;
float state[D];
#pragma unroll
for (int i = 0; i < D; i++) state[i] = 0.0f;
extern __shared__ float smem[];
float* k_smem = smem;
float* q_smem = smem + D;
for (int t = 0; t < T; t++) {
int base = ((b * T + t) * H + h) * D;
k_smem[j] = k[base + j];
q_smem[j] = q[base + j];
float gate_j = gate[base + j];
float v_j = v[base + j];
__syncthreads();
#pragma unroll
for (int i = 0; i < D; i++) {
state[i] = gate_j * state[i] + k_smem[i] * v_j;
}
float dot = 0.0f;
#pragma unroll
for (int i = 0; i < D; i++) {
dot += q_smem[i] * state[i];
}
out[base + j] = dot;
__syncthreads();
}
}
__global__ void gated_delta_attention_kernel_fallback(
const float* __restrict__ q,
const float* __restrict__ k,
const float* __restrict__ v,
const float* __restrict__ gate,
float* __restrict__ out,
int B,
int T,
int H,
int D
) {
int bh = blockIdx.x;
int b = bh / H;
int h = bh % H;
int j = threadIdx.x;
extern __shared__ float smem[];
float* k_smem = smem;
float* q_smem = smem + D;
float* state = smem + 2 * D + j * D;
for (int i = 0; i < D; i++) state[i] = 0.0f;
for (int t = 0; t < T; t++) {
int base = ((b * T + t) * H + h) * D;
k_smem[j] = k[base + j];
q_smem[j] = q[base + j];
float gate_j = gate[base + j];
float v_j = v[base + j];
__syncthreads();
for (int i = 0; i < D; i++) {
state[i] = gate_j * state[i] + k_smem[i] * v_j;
}
float dot = 0.0f;
for (int i = 0; i < D; i++) {
dot += q_smem[i] * state[i];
}
out[base + j] = dot;
__syncthreads();
}
}
```
For the fallback, the shared memory size is `(2*D + D*D) * sizeof(float)`.
Then in the wrapper:
```cpp
#define LAUNCH_REG(D) gated_delta_attention_kernel_reg<D><<<num_blocks, D, 2*D*sizeof(float)>>>( \
q_ptr, k_ptr, v_ptr, gate_ptr, out_ptr, B, T, H)
if (D == 24) {
LAUNCH_REG(24);
} else {
gated_delta_attention_kernel_fallback<<<num_blocks, D, (2*D + D*D)*sizeof(float)>>>(
q_ptr, k_ptr, v_ptr, gate_ptr, out_ptr, B, T, H, D);
}
```
This approach handles arbitrary D while giving D=24 (our test case) the best performance with register storage.
For the register version, since D is a template parameter, the compiler will likely inline all the `state[i]` accesses as 24 separate registers. This avoids local memory and bank conflicts. The shared memory for k_smem and q_smem is only 2*D floats, and the loads are coalesced.
Actually, I want to double check: D as a template parameter vs `float state[24]`. If we have `template <int D>` and instantiate with D=24, the compiler knows D=24 at compile time. In the loop `for (int i = 0; i < D; i++)`, with `#pragma unroll`, the compiler can fully unroll into 24 statements. And `state[i]` accesses in unrolled code might still be array accesses. But with optimization, the compiler promotes array elements to registers when the access pattern is simple.
In practice, nvcc is quite good at keeping small constant-sized arrays in registers. So `float state[24]` in a fully unrolled loop will likely use 24 registers.
But wait, there's still the issue of `state[i] = gate_j * state[i] + k_smem[i] * v_j`. In the unrolled version:
```
state[0] = gate_j * state[0] + k_smem[0] * v_j;
state[1] = gate_j * state[1] + k_smem[1] * v_j;
...
state[23] = gate_j * state[23] + k_smem[23] * v_j;
```
This has no loop-carried dependencies and can be executed with full instruction-level parallelism. Each state[i] is independent. The compiler will assign them to 24 different registers.
This is excellent.
For the dot product:
```
dot = q_smem[0] * state[0] + q_smem[1] * state[1] + ... + q_smem[23] * state[23];
```
This is a reduction tree. With 24 terms, it can be efficiently computed. nvcc might use FMAs.
So the register-based approach with template is very good.
One concern: `#pragma unroll` with a template parameter. I believe this works in nvcc. But to be safe, I can also not use `#pragma unroll` and let the compiler decide. For D=24, the compiler will likely unroll the loop automatically if optimizations are on.
Let me finalize the code:
```python
import torch
from torch.utils.cpp_extension import load_inline
cuda_source = """
#include <torch/extension.h>
#include <cuda_runtime.h>
template <int D>
__global__ void gated_delta_attention_kernel_reg(
const float* __restrict__ q,
const float* __restrict__ k,
const float* __restrict__ v,
const float* __restrict__ gate,
float* __restrict__ out,
int B,
int T,
int H
) {
int bh = blockIdx.x;
int b = bh / H;
int h = bh % H;
int j = threadIdx.x;
float state[D];
for (int i = 0; i < D; ++i) {
state[i] = 0.0f;
}
extern __shared__ float smem[];
float* k_smem = smem;
float* q_smem = smem + D;
for (int t = 0; t < T; ++t) {
int base = ((b * T + t) * H + h) * D;
k_smem[j] = k[base + j];
q_smem[j] = q[base + j];
const float gate_j = gate[base + j];
const float v_j = v[base + j];
__syncthreads();
for (int i = 0; i < D; ++i) {
state[i] = gate_j * state[i] + k_smem[i] * v_j;
}
float dot = 0.0f;
for (int i = 0; i < D; ++i) {
dot += q_smem[i] * state[i];
}
out[base + j] = dot;
__syncthreads();
}
}
__global__ void gated_delta_attention_kernel_fallback(
const float* __restrict__ q,
const float* __restrict__ k,
const float* __restrict__ v,
const float* __restrict__ gate,
float* __restrict__ out,
int B,
int T,
int H,
int D
) {
int bh = blockIdx.x;
int b = bh / H;
int h = bh % H;
int j = threadIdx.x;
extern __shared__ float smem[];
float* k_smem = smem;
float* q_smem = smem + D;
float* state = smem + 2 * D + j * D;
for (int i = 0; i < D; ++i) {
state[i] = 0.0f;
}
for (int t = 0; t < T; ++t) {
int base = ((b * T + t) * H + h) * D;
k_smem[j] = k[base + j];
q_smem[j] = q[base + j];
const float gate_j = gate[base + j];
const float v_j = v[base + j];
__syncthreads();
for (int i = 0; i < D; ++i) {
state[i] = gate_j * state[i] + k_smem[i] * v_j;
}
float dot = 0.0f;
for (int i = 0; i < D; ++i) {
dot += q_smem[i] * state[i];
}
out[base + j] = dot;
__syncthreads();
}
}
torch::Tensor gated_delta_attention_forward(
torch::Tensor q,
torch::Tensor k,
torch::Tensor v,
torch::Tensor gate
) {
int B = q.size(0);
int T = q.size(1);
int H = q.size(2);
int D = q.size(3);
auto out = torch::empty_like(q);
int num_blocks = B * H;
float* q_ptr = q.data_ptr<float>();
float* k_ptr = k.data_ptr<float>();
float* v_ptr = v.data_ptr<float>();
float* gate_ptr = gate.data_ptr<float>();
float* out_ptr = out.data_ptr<float>();
if (D == 24) {
gated_delta_attention_kernel_reg<24><<<num_blocks, 24, 2 * 24 * sizeof(float)>>>(
q_ptr, k_ptr, v_ptr, gate_ptr, out_ptr, B, T, H);
} else {
size_t smem_size = (2 * D + D * D) * sizeof(float);
gated_delta_attention_kernel_fallback<<<num_blocks, D, smem_size>>>(
q_ptr, k_ptr, v_ptr, gate_ptr, out_ptr, B, T, H, D);
}
return out;
}
"""
gated_delta_attention = load_inline(
name="gated_delta_attention",
cpp_sources="",
cuda_sources=cuda_source,
functions=["gated_delta_attention_forward"],
extra_cuda_cflags=["-O3"],
)
class ModelNew(torch.nn.Module):
def __init__(self):
super().__init__()
def forward(self, q, k, v, gate):
return gated_delta_attention.gated_delta_attention_forward(q, k, v, gate)
```
Wait, but the `state` array with `float state[D]` where D is a template parameter. Is this valid CUDA C++? Yes, in C++ you can have arrays with const size, and template parameters are compile-time constants.
However, to get the best registerization and unrolling, I should add `#pragma unroll` to the inner loops. But since the loops bounds are compile-time constants for the template case, nvcc will unroll them automatically at -O3.
Let me test with compile_kernel first.
Actually, looking at the compile tool, it takes the full Python kernel file code. So I need to wrap everything in a Python string and use `compile_kernel`.
Let me write the complete source and compile it.