reasoning
This is a Mamba S6 selective scan layer. Let me analyze the computation:
1. Input x: [B, L, D]
2. delta = softplus(proj_delta(x)) -> [B, L, D]
3. B_t = proj_B(x) -> [B, L, N]
4. C_t = proj_C(x) -> [B, L, N]
5. A = -exp(A_log) -> [D, N]
6. dA = exp(delta.unsqueeze(-1) * A.unsqueeze(0).unsqueeze(0)) -> [B, L, D, N]
7. dB = delta.unsqueeze(-1) * B_t.unsqueeze(2) -> [B, L, D, N]
8. x_db = x.unsqueeze(-1) * dB -> [B, L, D, N]
9. Loop t from 0 to L-1:
h = dA[:, t] * h + x_db[:, t]
y_t = (h * C_t[:, t].unsqueeze(1)).sum(dim=-1)
outputs.append(y_t)
10. y = stack(outputs) -> [B, L, D]
11. y = y + x * D
This is a scan/recurrence operation. The key bottleneck is the sequential dependency in the recurrence.
For a fused CUDA kernel, I should:
1. Compute delta, B_t, C_t in the kernel (or precompute them)
2. Perform the scan over L dimension
3. Apply D skip connection
Since L=2048 and D=256, N=16, B=16:
- The inner dimension per scan is D*N = 256*16 = 4096
- Total elements per batch: 2048 * 256 * 16 = 8.3M per batch, *16 = 133M elements to scan
For GPU optimization, I can:
1. Use thread-level parallelism over B and D
2. Each thread handles one (B, D) pair and scans over L, with N elements of state
3. This means each thread block handles a subset of (B, D) pairs
4. Within a thread, we process L sequentially but N in parallel (or vice versa)
Since N=16 is small, I could have each thread handle one (B, D) and process N=16 state elements.
Actually, let me think about the memory access pattern:
- For each (b, d), we need to read delta[b, :, d], B_t[b, :, :], C_t[b, :, :], A[d, :], x[b, :, d]
- Then compute the scan over t
- Output y[b, :, d]
For best performance, I should process multiple (B, D) pairs per block, with each thread handling one (b, d) pair and doing a loop over t, processing N=16 elements.
Wait, but we also need to compute delta, B_t, C_t from x. The projections are matmuls. For the custom kernel, I have two options:
1. Call the projections in PyTorch, then pass them to the kernel
2. Try to fuse the projections too
Option 1 is much simpler. Let me see if it's enough. The projections are:
- proj_delta: [B, L, D] @ [D, D] -> [B, L, D] (actually Linear(D, D) with bias)
- proj_B: [B, L, D] @ [D, N] -> [B, L, N]
- proj_C: [B, L, D] @ [D, N] -> [B, L, N]
These are standard matmuls that cuBLAS can do efficiently. fusing them into a custom kernel might be hard to beat cuBLAS on since D=256, N=16.
Actually wait - let's think about this more carefully. The projections + scan together is end-to-end. But the scan is the real bottleneck due to sequential dependency.
For the scan, each (b,d) is independent. For each (b,d):
- State h[N] starts at 0
- For t in 0..L-1:
h[n] = dA[t,n] * h[n] + x_db[t,n]
y[t] = sum_n h[n] * C[t,n]
Wait, dA depends on delta and A. dA = exp(delta * A). Since A is [D, N] and delta is per (b,l,d), dA is per (b,l,d,n).
Similarly x_db = x * delta * B.
So for each (b,d):
- Read A[d, 0..N-1] (constant for the scan)
- For t = 0..L-1:
- delta_t = delta[b,t,d]
- x_t = x[b,t,d]
- B_t_n = B[b,t,0..N-1]
- C_t_n = C[b,t,0..N-1]
- For n = 0..N-1:
- dA_tn = exp(delta_t * A[d,n])
- h_n = dA_tn * h_n + x_t * delta_t * B_t_n
- y_t = sum_n h_n * C_t_n
This is a per-channel scan. Each (b, d) is completely independent.
The parallelism is over B * D = 16 * 256 = 4096 scan instances. Each scan is over L=2048 steps, with inner dimension N=16.
4096 parallel scans is enough to saturate the GPU. Each thread can handle one (b, d) scan over L.
For an H100 with 132 SMs, we want at least 4096 thread blocks or more. If each block handles 1 (b,d), we get 4096 blocks which is plenty.
Actually, let me think about whether we should use 1 thread per (b,d) or a warp or block.
- N=16 is small
- L=2048 is the sequential dimension
- Each step requires: computing dA for N elements, updating h for N elements, computing dot product with C for N elements
- At fp32: about 16* several ops per step
- Total work per (b,d): ~2048 * ~50 ops = ~100K ops
A single thread doing this might be fine but we want memory coalescing. Since N=16, a warp of 32 threads could process 2 (b,d) pairs, with each thread handling all N=16 elements sequentially, or each thread handling half the N elements.
Actually for better memory bandwidth, having a warp process multiple (b,d) elements contiguous in D would help coalesce reads from delta[B, L, D], x[B, L, D].
Wait, the reads from B and C are [B, L, N], so for fixed (b,t), B_t is contiguous in N (if N is last dim). But we need B[b, t, :] for each t, which is contiguous.
Let me think about the layout. If we process (b,d) with one thread:
- delta reads: delta[b, :, d] - stride D
- x reads: x[b, :, d] - stride D
- B reads: B[b, t, :] - contiguous for each t
- C reads: C[b, t, :] - contiguous for each t
- A reads: A[d, :] - contiguous
These reads are strided for delta and x, but contiguous for B, C, A. With D=256, the delta/x strides are 256, which means adjacent threads reading adjacent d would have coalesced access.
So if thread t processes (b, d_t), and d_t varies across threads in a warp, reads from delta[b, t, d_t] are coalesced (consecutive d, stride 1 in the contiguous dimension if D is last).
Actually in PyTorch, [B, L, D] has stride [L*D, D, 1]. So delta[b, t, d] is at offset b*L*D + t*D + d. Adjacent d values are adjacent in memory. So if a warp reads delta[b, t, d] through delta[b, t, d+31], this is perfectly coalesced.
But B[b, t, :] is at offset b*L*N + t*N + n.
So one approach: each warp handles WARP_SIZE (32) different (b, d) channels. Each thread loops over L and over N.
For N=16, each thread would do 16 operations per L step. That might be okay.
Alternatively, we can have a block with 32 threads handle multiple (b,d) pairs, processing them sequentially.
Let me try a simpler approach first: one thread per (b,d). Since there are 4096 (b,d) pairs, and each SM can handle 2048 threads, this gives 2 SMs worth of work if fully packed. But actually each thread does a lot of work (L=2048 steps), so occupancy might be limited by register usage.
With 4096 threads across ~100 warps, on 132 SMs, each SM gets less than 1 warp. That's actually quite thin. We'd want more parallelism.
Wait, we can have each block with say 256 threads handle 256 (b,d) channels for one batch item. There are B=16 batches, so 16 blocks total. That's too few blocks - would severely underutilize the GPU.
Better: each block with T threads handles T channels spread across all batches. For example, 256 threads per block, each thread handles one (b,d). Then grid = ceil(4096/256) = 16 blocks. Still too few.
We need more blocks. So either:
1. Multiple blocks per "slice" - not straightforward for independent channels
2. Have each thread handle fewer steps - split L dimension among blocks/threads
Actually, we can split L into chunks, but the scan recurrence makes that harder (need prefix scan algorithm).
For a simple approach with enough parallelism:
- Each block = 128 threads
- Each thread = 1 (b,d) pair
- We get 4096/128 = 32 blocks... still not great.
What if we use a block size of 64? 4096/64 = 64 blocks. Still not great for 132 SMs.
Hmm, but actually the work per thread is quite high. 2048 iterations with ~20-30 FLOPs each = ~50K FLOPs. At 67 TFLOPS, a warp can do this in ~1 microsecond. 4096 threads = 128 warps. At full utilization, this is very fast.
The real question is memory bandwidth. Each (b,d) scan reads:
- delta[b, :, d]: L floats = 8KB
- x[b, :, d]: L floats = 8KB
- B[b, :, :]: L*N floats = 32KB for N=16, L=2048
- C[b, :, :]: L*N floats = 32KB
Total: ~80KB reads per scan
For 4096 scans: 80KB * 4096 = 327MB of reads
At 3.35 TB/s = 3350 GB/s, this would take 0.1ms for memory alone.
Plus writes: y[b, :, d]: L floats = 8KB * 4096 = 32MB, ~0.01ms.
So total time ~0.1ms just for memory. The compute is trivial.
Wait, but B and C reads: each scan instance (b,d) reads all of B[b, :, :] and C[b, :, :], which is independent of d! That's a huge waste. B and C don't depend on d.
For B=16, L=2048, N=16:
- B tensor size: 16 * 2048 * 16 = 524288 floats = 2MB
- Total B reads if each of 4096 scans reads all B: 4096 * 2MB / 16 = 512MB? Wait no, B is [B, L, N], each scan for batch b reads B[b, :, :]. For all d in batch b, they all read the same B[b, :, :].
Actually for B=16 batch items, each batch item has 256 d-channels. But B[b, :, :] is shared across all d channels in batch b. So total B reads = 16 * 2048 * 16 floats read 256 times each = not good for cache.
256 reads of the same data for 256 channels in the same batch. If they're running concurrently in the same SM or nearby SMs, the L2 cache might help. 2MB * 256 = 512MB of B reads total. Similarly for C.
But the cache line holding B[b, t, :] is only N*4=64 bytes. 256 threads in the same warp/block might share it well.
Actually, for adjacent d channels in the same batch b, they all read identical B[b, t, :] and C[b, t, :]. So if we group threads by (b, d) adjacency in d, they'll have the same B/C access pattern and can share cache lines.
This is actually very cache-friendly if threads in the same warp process the same batch b and nearby d.
OK so here's my approach:
- Grid: enough blocks to cover all (B, D) = 4096 elements
- Block size: say 256 threads (or 128)
- Thread i in block processes (b, d) = (blockIdx.x * blockDim.x + threadIdx.x) % (B*D)
- b = idx / D
- d = idx % D
- Loop over t = 0..L-1:
- Read delta, x for this (b, d, t)
- Read B[b, t, :] and C[b, t, :]
- Update h[0..N-1] using A[d, :]
- Compute y_t = dot(h, C_t)
- Write y[b, t, d]
Since N=16, I can keep h in registers. dA and dB computations are per step.
Wait, but B[b, t, :] are the same for all d in batch b. So in a block, if threads handle different d in the same b, they all read the same B/C values. We should leverage this by having warps within a block process the same b but different d.
For example, a block of 256 threads could process d=0..255 for a single batch b. Then all threads in the block read the same B[b, :, :] and C[b, :, :] at each step, which means excellent cache sharing.
But we have B=16 batch items. If each block handles one batch item, we only get 16 blocks. To get more parallelism, we can have multiple blocks per batch item.
For batch b, we could have block_count blocks, each handling contiguous d ranges. Or we can have blocks handle L chunks with a more complex scan.
Alternatively, process multiple L steps per thread with unrolling, and use multiple blocks per (b, d). But the recurrence prevents easy splitting.
Actually for 16 blocks, with 256 threads each, on 132 SMs - H100 can run multiple blocks per SM. Even 16 blocks should be fine if each block has enough occupancy and work. At L=2048 steps per thread with compute, a block of 256 threads has 256*2048 = 524k iterations. That's a lot of work. But will 16 blocks saturate the GPU?
Hmm, with 132 SMs and only 16 blocks, most SMs will be idle. On H100, each SM can run at most a certain number of blocks concurrently. If each block uses a lot of registers, maybe only 1-2 blocks per SM. So 16 blocks might use ~16 SMs, leaving 116 SMs idle.
To fix this, we need to increase parallelism. Options:
1. Have K blocks per batch item, each handling D/K channels
2. Or split L: have each block handle a subsequence, but this requires cross-block communication for the scan (Blelloch scan or similar)
For option 1, if we use 8 blocks per batch item, we get 128 blocks total. Each block handles 256/8 = 32 d channels (32 threads). We can increase to 16 blocks per batch item = 256 blocks, each handling 16 d channels with say 64 threads (one thread per d, but 64 threads per block means each thread processes 0.25 d channels... doesn't work well).
Actually, we could just launch many more threads per d channel. But each d channel is a single scalar scan. Hard to parallelize the scalar scan itself.
Unless... we parallelize over N. N=16 is small. But we could have a warp (32 threads) process one (b,d) scan, with each thread handling n = threadIdx.x % 16 or similar. But then we'd only have 4096/32 = 128 warps across all blocks, which is thin.
Let me think about option 2 more carefully: splitting L.
For an associative scan / prefix scan, we can divide L into chunks. But is the recurrence associative?
h_t = dA_t * h_{t-1} + x_db_t
This is of the form:
h_t = a_t * h_{t-1} + b_t
This can be written as an associative scan using linear functions:
Let f_t(h) = a_t * h + b_t
Then f_s ∘ f_t (h) = a_s * (a_t * h + b_t) + b_s = (a_s * a_t) * h + (a_s * b_t + b_s)
So each step can be represented as a pair (a, b), and composition is:
(a2, b2) ∘ (a1, b1) = (a2*a1, a2*b1 + b2)
For our case with vector h[N]:
For each n: h_t[n] = dA_t[n] * h_{t-1}[n] + x_db_t[n]
So for each n, we have a scalar affine recurrence. The N dimensions are independent!
So we can process each n independently with a prefix scan. But N=16 is small.
Actually, this means each (b, d, n) is an independent scalar affine scan. That's B*D*N = 16*256*16 = 65536 independent scalar scans!
65536 parallel scans, each over L=2048. Now that's plenty of parallelism.
For each (b, d, n):
- a_t = dA[b, t, d, n] = exp(delta[b, t, d] * A[d, n])
- b_t = x_db[b, t, d, n] = x[b, t, d] * delta[b, t, d] * B[b, t, n]
- h_t = a_t * h_{t-1} + b_t
- Contribution to y[b, t, d] = h_t * C[b, t, n]
So y[b, t, d] = sum_n h_t[b, d, n] * C[b, t, n]
If we parallelize over (b, d, n), we need to sum over n to get y. That's a reduction per (b, t, d).
Hmm, that adds complexity. We'd need to compute all h_t for all n, then sum across n.
With 65536 parallel instances and L=2048, each instance does 2048 steps of the affine recurrence. After computing h_t, we need to do y[b, t, d] += h_t * C[b, t, n].
To handle the reduction, we could:
1. Compute h for all (b, d, n, t) and store it -> huge memory: 16*256*16*2048 = ~134M floats = 536MB
2. Or compute in cooperative groups: each (b, d) has a group of N=16 threads, each processing one n, and they do the scan while a leader thread or warp-shuffle accumulates y.
Actually, with N=16, a half-warp or a warp could process all n for one (b, d). Let me think.
For each (b, d):
- A warp of 16 threads (or 32, with 2 per n, or just 16 active threads) processes n=0..15
- Each thread loops over t=0..L-1:
- Read delta_t, x_t (broadcast to all threads in warp, or read once per warp)
- Read B_t_n, C_t_n for its n
- Compute a_t = exp(delta_t * A_n)
- Compute b_t = x_t * delta_t * B_t_n
- Update h_n = a_t * h_n + b_t
- Compute contribution = h_n * C_t_n
- Warp-reduce (sum) contributions to get y_t
- Thread 0 writes y_t
Wait, warp-level reduction at each step. That's L * warp_reduce operations per (b, d). A warp reduce of 16 values is cheap (a few shuffles).
So per (b, d):
- 1 warp (or 16 threads out of 32)
- Each warp handles 1 (b, d)
- Total warps needed: 4096 warps
- Threads needed: 4096 * 32 = 131072 threads (if using full warp per (b,d)) or 4096 * 16 = 65536 threads
Using full warp (32 threads) per (b,d) with 16 threads actively doing work:
- 4096 warps
- On H100, each SM can hold 64 warps (with 2048 threads max), so total SMs needed if fully packed: 4096/64 = 64 SMs. With 132 SMs, that's about half utilization, but we can use 2 blocks per SM.
Actually, we can have blocks of say 4 warps = 128 threads, where each warp handles a different (b, d). 4096 warps / 4 = 1024 blocks. That's plenty!
Block size = 128 threads (4 warps). Grid = 1024 blocks (for 4096 warps total, 4 warps per block). Wait no, if each warp handles 1 (b,d), and we have 4096 (b,d) pairs, we need 4096 warps. With 4 warps per block, that's 1024 blocks. 1024 blocks on 132 SMs = ~8 blocks per SM, which is great for occupancy.
If we instead use 16 threads per (b,d) (half warp), we could pack 2 (b,d) per warp, giving 2048 warps total. With 4 warps per block = 512 blocks. Still decent.
But half-warps aren't natural in CUDA (warp size is 32). We could just use 16 threads in a warp and leave 16 idle, or map 2 n-values per thread.
Actually, a cleaner approach: use 1 thread per (b, d, n)! That's 65536 threads. For the reduction to get y, we need to sum across n for each (b, d, t).
If 16 threads (one per n) work together for one (b, d), they can do warp-shuffle reductions. That's the same as 16 threads per (b,d).
So: 1 block = 256 threads = 16 warps. Each warp handles 2 (b, d) pairs? No, if we use 16 threads per (b,d), then a warp of 32 threads handles 2 (b,d) pairs. But 32 doesn't divide evenly by 16? It does: 2 pairs per warp.
So 256 threads / 32 threads per warp = 8 warps. Each warp = 2 (b,d) if we use 16 threads per pair. 8 * 2 = 16 (b,d) pairs per block.
Total blocks = 4096 / 16 = 256 blocks. 256 blocks / 132 SMs = ~2 blocks per SM. That's workable but might be a bit low.
Alternatively, 128 threads per block (4 warps), each warp handles 2 (b,d), total 8 (b,d) per block, total 512 blocks. Better.
Or 256 threads per block but mapping differently: 256 threads = 16 (b,d) with 16 threads each.
Let me do the math:
- 65536 total (b,d,n) elements
- Block size 256 threads
- Threads per (b,d) = 16 (one per n)
- (b,d) per block = 16
- Blocks needed = 4096 / 16 = 256 blocks
With 256 blocks and 132 SMs, most SMs get 1-2 blocks. But each block has 16 warps, and each warp does 2 (b,d) using 16 threads each. The occupancy might be limited by registers.
Actually, a warp with 16 active and 16 inactive threads is fine, but the scheduler still sees 32 threads.
Let me think about registers. Each thread:
- h_n (1 float)
- A_n (1 float)
- Some temporaries
Total maybe 8-10 registers per thread. At 256 threads per block, that's well within limits.
OK, but there's another issue. If we have 256 blocks, that's not enough to fully utilize 132 SMs if each SM can run only a few blocks concurrently. Actually, on H100, each SM can run up to 32 blocks. With 256 blocks total and 132 SMs, that's 256/132 ≈ 2 blocks per SM, which is fine.
But wait, can we do better? What if we split L into chunks and have multiple blocks per (b,d) working on different L segments?
For each (b,d) pair, the scan is sequential. We'd need a prefix-scan algorithm for L. For long L, this might help.
Actually L=2048 isn't that long. With 256 blocks and each thread doing 2048 steps, total thread-steps = 65536 * 2048 = 134M steps. At say 10 ops per step = 1.34B ops. The GPU can do 67 TFLOPS, so compute is ~0.02ms. Memory is the bottleneck.
For memory, each of 65536 threads reads:
- delta: shared among n in same (b,d). So 16 threads per (b,d) read the same delta. Need to ensure good caching or explicit shared memory loading.
- x: also shared.
- B_t_n and C_t_n: each thread reads one n, so together they read all N=16 values. This is contiguous if B[b,t,:] is contiguous.
- A_n: constant per thread.
Let's use shared memory for delta and x. But with 16 threads per (b,d) and multiple (b,d) per block, we can load delta and x into shared memory for all (b,d) in the block.
For block with 16 (b,d) pairs:
- L=2048 steps
- At each step t, we need delta[b, t, d_i] and x[b, t, d_i] for i=0..15
- These are 16 floats each.
- The block loops over t, loads these into shared memory (or registers), then all threads use them.
Actually, for 16 (b,d) pairs in a block, at step t:
- delta values: delta[b_i, t, d_i] for i=0..15
- x values: similarly
- These are scattered in memory because (b_i, d_i) might not be contiguous.
If we assign (b,d) contiguously to threads in a block, then for threads in same warp, adjacent threads have adjacent d (same b). Then delta[b, t, d] and delta[b, t, d+1] are adjacent in memory, so coalesced.
Good!
Block assignment:
- block_idx handles (b,d) from start_idx = block_idx * pairs_per_block to start_idx + pairs_per_block - 1
- For pair i in 0..pairs_per_block-1, thread_base = i * 16, handles n=0..15
- (b,d) = index / D, index % D, where index = start_idx + i
For coalescing, if pairs_per_block = 16, and the 16 pairs are contiguous in the flat (b*D + d) space, then the 16 pairs have either:
- Same b, d=0..15 (if at start of b), or
- Cross boundary between b and b+1
If same b and contiguous d, then delta[b, t, d: d+16] is contiguous in memory, perfect coalescing.
Within a warp of 32 threads:
- If we have pairs_per_block = 16, with 16 threads each, that's 256 threads total = 8 warps.
- Warp 0: pair 0, n=0..15 (threads 0..15), pair 1, n=0..15 (threads 16..31)?
- Actually 32 threads: first 16 for pair 0, next 16 for pair 1.
In warp 0, threads read:
- t0..t15: delta[b, t, d] (same b, t, d), A[d, n] (vary in n), ...
- t16..t31: delta[b', t, d'] (same b', t, d'), A[d', n], ...
If pair 0 and pair 1 are adjacent (d+1 or next b), then delta[b, t, d] and delta[b', t, d'] are adjacent or nearby. Should be coalesced.
For B and C reads, all 32 threads in warp 0 read B[b, t, n] and B[b', t, n] respectively, with varying n.
- Thread i (0..15) reads B[b, t, i]
- Thread j (16..31) reads B[b', t, j-16]
If b != b', these are in different memory regions, so two cache lines. Each group of 16 threads reads 16 contiguous floats = 64 bytes = 1 cache line. Perfect.
This looks great!
Let me formalize the design:
**Grid**: 256 blocks (to cover 4096 (b,d) pairs with 16 per block)
**Block**: 256 threads, organized as 16 pairs × 16 threads per pair
**Shared memory**: Temporary buffers for delta, x, or just rely on L2 cache
Actually, let me not overcomplicate with shared memory for delta/x. Each of the 16 threads in a pair reads the same delta and x values. If they're in the same warp, broadcast from L1 cache handles this. If they're in different warps, L2 cache should help. Since we have 256 threads and 16 pairs, only 16 unique delta values and 16 unique x values per step. That's very small.
Wait, 16 threads in a pair reading the same delta value: this is a broadcast within the warp. The hardware handles this efficiently with a single cache line read.
For B and C: 16 threads in a pair read 16 contiguous floats (B[b, t, 0..15]), also 1 cache line.
So memory bandwidth should be near optimal.
Let me think about the exact kernel structure:
```cuda
__global__ void selective_scan_kernel(
const float* x, // [B, L, D]
const float* delta, // [B, L, D]
const float* A_log, // [D, N]
const float* B, // [B, L, N]
const float* C, // [B, L, N]
const float* D_skip, // [D]
float* y, // [B, L, D]
int B, int L, int D, int N
) {
int total_pairs = B * D;
int pairs_per_block = blockDim.x / N; // 256/16 = 16
int pair_idx_base = blockIdx.x * pairs_per_block;
int tid = threadIdx.x;
int pair_idx_in_block = tid / N; // 0..15
int n = tid % N; // 0..15
int pair_idx = pair_idx_base + pair_idx_in_block;
if (pair_idx >= total_pairs) return;
int b = pair_idx / D;
int d = pair_idx % D;
// Load A
float A_val = -expf(A_log[d * N + n]);
// State
float h = 0.0f;
// Pointers
const float* x_ptr = x + b * L * D + d;
const float* delta_ptr = delta + b * L * D + d;
const float* B_ptr = B + b * L * N;
const float* C_ptr = C + b * L * N;
float* y_ptr = y + b * L * D + d;
for (int t = 0; t < L; t++) {
float delta_t = delta_ptr[t * D];
float x_t = x_ptr[t * D];
float B_tn = B_ptr[t * N + n];
float C_tn = C_ptr[t * N + n];
float dA = expf(delta_t * A_val);
float dB = delta_t * B_tn;
float x_db = x_t * dB;
h = dA * h + x_db;
float contrib = h * C_tn;
// Warp reduce sum across n=0..15
// We need to organize so that the 16 threads for this pair can shuffle
int lane = tid % 32;
// If pair_idx_in_block is even, this pair uses lanes 0..15
// If odd, lanes 16..31
int lane_base = (pair_idx_in_block % 2) * 16;
// Put contribution in correct position for reduction
// Actually, easier: all 16 threads for a pair are contiguous in lane
// (e.g., lanes 0-15 or 16-31)
// Warp shuffle sum within the 16 lanes
// But warp shuffle works on the whole warp of 32.
// If only 16 threads have valid data in the 32-thread warp...
// Actually, since pair0 uses lanes 0-15 and pair1 uses lanes 16-31,
// we can do a segmented reduction with masks, or we can restructure.
}
}
```
Wait, the warp shuffle is tricky with 16-thread subgroups in a 32-thread warp. We can use `__shfl_xor_sync` with masks, or just use the `_sync` primitive with masks.
Actually, `__shfl_down_sync` with mask is the way to go. But for simplicity, let me organize threads differently.
Option: Use 16 threads per pair, but make each warp contain 2 pairs. Then for warp shuffle, we can use masks:
- Pair in lanes 0-15: mask = 0x0000FFFF
- Pair in lanes 16-31: mask = 0xFFFF0000
But `__shfl_down_sync` syntax:
```cuda
float __shfl_down_sync(unsigned mask, float var, unsigned int delta, int width=warpSize);
```
With width=16, it does reduction within sub-groups of 16 threads. But we need to make sure the 16 threads that form a pair are in the same sub-group.
In a 32-thread warp, threads are divided into sub-groups of `width`. If width=16, thread 0-15 form one group and thread 16-31 form another group. This is perfect if pair0 is in 0-15 and pair1 is in 16-31!
So:
```cuda
float sum = __shfl_down_sync(0xFFFFFFFF, contrib, 8, 16);
sum += __shfl_down_sync(0xFFFFFFFF, sum, 4, 16);
sum += __shfl_down_sync(0xFFFFFFFF, sum, 2, 16);
sum += __shfl_down_sync(0xFFFFFFFF, sum, 1, 16);
```
After this, thread 0 of each 16-thread sub-group has the sum. It writes to y.
But wait, there's a subtlety. `__shfl_down_sync` with width=16 partitions the warp into two 16-thread partitions. Thread X in the first partition only communicates with threads 0..15, and thread Y in second partition only with threads 16..31. This is exactly what we want!
However, thread 0 and thread 16 need to be the "leaders" of their respective pairs.
In my layout:
- tid = pair_idx_in_block * N + n
- For pair 0: tid in [0, 15]
- For pair 1: tid in [16, 31]
- For pair 2: tid in [32, 47]
- etc.
Wait no, pairs_per_block = 16, so tid = 0..255.
- Pair 0: tid 0..15
- Pair 1: tid 16..31
- Pair 2: tid 32..47
- Pair 3: tid 48..63
...
In warp 0 (tid 0..31): pairs 0 and 1.
In warp 1 (tid 32..63): pairs 2 and 3.
...
Perfect! Each warp has exactly 2 pairs, and the 16 threads of each pair are in the correct sub-group for width=16 shuffle.
After the warp shuffle, thread leader (n=0) of each pair has the sum. In the warp, these are tid=0 and tid=16.
So:
```cuda
if (n == 0) {
y_ptr[t * D] = sum + x_t * D_skip[d];
}
```
Wait, only the leader has access to the sum, and only the leader writes. But other threads also need to continue the recurrence with their own h value, which they have. That's fine; only the sum needs to be reduced and written.
One detail: the mask `0xFFFFFFFF` in `__shfl_down_sync` is correct because all 32 threads in the warp are active.
Let me verify the shuffle logic:
For a thread with lane `lane_id = tid % 32`:
- If tid is in 0..15, lane_id is 0..15. In the sub-group of width=16, its index is lane_id.
- `__shfl_down_sync(..., 8, 16)` gets the value from thread with index lane_id + 8 within the same sub-group.
- Correct!
So the reduction will give pair 0's sum to thread 0, and pair 1's sum to thread 16.
Great, this works.
Now let me also think about whether we want to use vector loads or anything. N=16 is a reasonable size. B, C have N as the last dimension, so B[b,t,0..15] is contiguous. We could use float4 loads to read 4 floats at a time. But B and C are used one element at a time in the loop, and the values are needed in natural order. Vector loading might help if we process N in groups.
Actually, the loop processes t sequentially. At each t, we read B_t[n] and C_t[n]. Since n is thread-specific, each thread reads its own offset. No vectorization possible here unless we reorganize.
But wait, we could restructure: instead of each thread handling one n and looping over t, each thread could handle multiple n and still loop over t. Since N=16, and a warp has 32 threads, we could have each warp handle 2 (b,d) pairs with each thread handling one n. Or we could have each warp handle 1 (b,d) pair with 2 threads per n to increase parallelism. But with 2 threads per n, they'd do redundant work except for the shuffle reduction being distributed differently.
Actually, having each warp handle 1 (b,d) with 32 threads, where n=0..15 are handled by 2 threads each (redundant), or where 16 threads are active and 16 idle, seems wasteful. The 2-pairs-per-warp approach is better.
Another optimization: unroll the t loop. L=2048. We can unroll by say 4 or 8. But 2048 is not necessarily divisible by some power of 2? It is: 2048 = 2^11. Unrolling by 4 or 8 should be fine.
However, unrolling a 2048-iteration loop might increase code size. Let's leave it to the compiler or unroll manually in chunks of 4.
Actually, for H100 throughput, manual unrolling might help hide latency, especially for expf and memory loads.
Let me think about instruction-level parallelism. Each iteration has:
1. Load delta_t, x_t
2. Load B_tn, C_tn
3. Compute dA = expf(delta_t * A)
4. Compute dB = delta_t * B_tn
5. Compute x_db = x_t * dB
6. Update h = dA * h + x_db
7. Compute contrib = h * C_tn
8. Shuffle reduce
9. Write
Steps 3-7 have dependencies. But expf is relatively slow (~10-20 cycles). We can overlap iterations by processing 2 t-values at a time if we have enough registers. But the recurrence h_t depends on h_{t-1}, so we can't overlap much.
Wait, within one thread, h_{t+1} = dA_{t+1} * h_t + x_db_{t+1}. The expf for t+1 can start while we're computing h_t * C_tn. There's some ILP.
Let me not overcomplicate. The compiler handles this reasonably well.
Another idea: precompute A_val = -expf(A_log[d*N+n]). A_log is in global memory. That's fine.
Should we fuse the projections into the kernel?
If we precompute delta, B_t, C_t in PyTorch, we need to:
1. Allocate delta, B_t, C_t
2. Run 3 matmuls + softplus
3. Pass to kernel
4. Kernel computes scan
The matmuls are:
- proj_delta: [16, 2048, 256] @ [256, 256] = ~2.1M output, ~0.5B MACs
- proj_B: [16, 2048, 256] @ [256, 16] = ~134M MACs
- proj_C: same as proj_B
These are not tiny. [16, 2048] is 32768 tokens. For proj_B with 256->16: this is 32768 * 256 * 16 = 134M MACs. At 989 TFLOPS (TF32 tensor core), this is ~0.14ms. Actually wait, this is matrix-vector in batch mode.
PyTorch linear will call cuBLAS. For 32768 tokens and D=256, N=16, it's roughly a GEMM of M=32768, K=256, N=16. cuBLAS is highly optimized for this. Fusing into our custom kernel might not be easy to beat.
However, the scan itself with naive PyTorch loop is likely much slower. Let me check the reference: it materializes dA [B, L, D, N] = 16*2048*256*16 = ~134M floats = 536MB, and x_db same size. Then does a Python loop over L. That's incredibly slow.
So if we just replace the scan part with a fused CUDA kernel, we should already see huge speedup. The projections can remain in PyTorch.
Wait, but we can be even smarter. If we don't precompute delta, B_t, C_t, we avoid writing them to global memory. But PyTorch's linear layers will write their outputs to global memory anyway. We can pass the input x to our kernel and compute proj_B, proj_C, proj_delta inside... no, we don't have the weight matrices as simple constants, they're parameters of the model.
The kernel is called from ModelNew.forward. It has access to self.proj_delta.weight, self.proj_B.weight, etc. We could either:
1. Run projections in PyTorch, pass results to kernel
2. Pass weights to kernel and do the matmuls inside
Option 2 is very complex to implement from scratch and unlikely to beat cuBLAS for these sizes. I'll do option 1.
Actually, we can optimize the memory layout. The kernel reads x, delta, B, C, A_log, D and writes y. If PyTorch computes delta, B, C as separate tensors, they live in global memory and our kernel reads them.
Can we avoid the separate allocations? We can compute delta, B, C inside the kernel on-the-fly if we pass the weights and do the matmuls. But again, hard to beat cuBLAS.
Wait, let me reconsider. For proj_B: it's x @ W_B where W_B is [D, N]. For a single token, this is dot(x[t,:], W_B[:,n]). If we're already looping over t in the kernel, and N=16, then for each t we need to compute B_t[n] = dot(x[t,:], W_B[:,n]).
This is N dot products of length D=256. At each step t, we'd do 256*16 = 4096 multiplies and 4096 adds. Over L=2048, that's 8M MACs per (b). For B=16, 134M MACs total. cuBLAS does this efficiently with tensor cores.
If we do it in our kernel, each thread block would need to do the matmul for its batch item. But our blocks only handle 16 d-channels out of 256. The B and C projections output N channels that don't depend on d, so all d-channels in a batch need the same B and C values.
If we fuse the projections, we'd compute B and C on-demand or precompute them in shared memory before the scan. For a block handling one batch item and some d-channels, it needs B and C for all t. That's the entire B[b, :, :] and C[b, :, :] tensors: 2048*16 = 32k floats each = 128KB. Shared memory per block on H100 is 227KB max. We could fit both.
But all threads in the block (256 threads) can cooperatively compute B and C for one t, or precompute all t before the scan. Precomputing all 2048*16 = 32k values in shared memory is 128KB. Each thread would then read from shared memory during the scan.
Hmm, but computing B in the kernel: for each t, we need B[b, t, n] = sum_d x[b, t, d] * W_B[d, n]. This is a matvec. The block could have all threads compute this cooperatively. But we have 256 threads for 16 d-channels. If the block handles d=0..255 for one batch, then all threads work on different d, not the same t. Actually no, each block handles 16 (b,d) pairs, not all 256 d-channels for one b.
Wait, if each block handles 16 pairs, and there are 16 pairs per block with 16 threads each, the pairs could be 16 d-channels from one batch, or spread across batches. To get maximum coalescing, let's say each block handles 16 contiguous d for one batch. Then b is the same for all pairs in the block.
In this case, B[b, t, :] and C[b, t, :] are the same for all pairs in the block. They could be computed cooperatively by the block and stored in shared memory.
But computing B[b, t, n] for n=0..15 requires x[b, t, 0..255] and W_B[0..255, 0..15]. With 256 threads, each thread could compute one n for all d? No, there are 256 threads and only 16 n values. Or each thread computes a partial sum for one or more n.
Actually, 256 threads could each compute one element of x[b,t,:] * W_B[:, n] for specific (d, n) combinations, then do a block-level reduction. But that's very complex and might not be faster than just computing the matvec with a few threads or just precomputing B,C in PyTorch.
Given the complexity and that PyTorch+cuBLAS handles the projections efficiently, I'll stick with precomputed B, C, delta.
So the ModelNew.forward does:
1. delta = F.softplus(self.proj_delta(x))
2. B_t = self.proj_B(x)
3. C_t = self.proj_C(x)
4. Call kernel(x, delta, B_t, C_t, self.A_log, self.D, y)
5. Return y
The kernel then does:
- Compute A = -exp(A_log)
- For each (b, d), scan over t
- At each t, compute dA, dB, update h, compute contribution, reduce, write y
Let me write the kernel more carefully.
Wait, I need to think about numerical precision. The reference computes A = -exp(A_log) once in forward. Then dA = exp(delta.unsqueeze(-1) * A.unsqueeze(0).unsqueeze(0)).
In the kernel, I compute A_val = -expf(A_log[d*N + n]). Then dA = expf(delta_t * A_val). This matches exactly.
dB in reference: delta.unsqueeze(-1) * B_t.unsqueeze(2) -> for fixed (b,t,d), dB is delta[b,t,d] * B[b,t,:]. My kernel computes dB = delta_t * B_tn. Matches.
x_db in reference: x.unsqueeze(-1) * dB -> x[b,t,d] * dB. My kernel: x_t * dB. Matches.
Then h = dA * h + x_db. Matches.
y_t = (h * C_t[:,t].unsqueeze(1)).sum(dim=-1).
In kernel: h * C_tn, then sum over n. Matches.
Finally y = y + x * D. Matches.
Good.
Now let me think about whether I should also handle the projections in a more fused way.
Actually, there is one bandwidth optimization: PyTorch will write delta, B, C to global memory, then the kernel reads them. That's 3 * [B,L,D or N] written and read. But PyTorch might fuse some of these computations within its own kernels.
For delta: softplus(proj_delta(x)) - PyTorch likely does this as matmul + softplus elementwise. The elementwise softplus could in theory be fused with the matmul but I'm not sure PyTorch does that.
Actually, we can do better. Instead of having the kernel read delta from global memory, we can have it read x and compute delta = softplus(x @ W_delta + b_delta) inside the kernel. But again, this requires a matmul.
For proj_delta: x[b,t,:] @ W_delta[:,d] + b_delta[d] for each (b,t,d). This is very similar to what we have for B and C. But delta is needed per (b,t,d), which is exactly what our kernel processes. Each thread needs delta[b,t,d] for its specific (b,d) and all t.
If the block handles 16 d-channels for one batch, it needs x[b,t,0..255] for all t, and W_delta[0..255, d_i] + b[d_i] for each of the 16 d-channels.
Actually, delta[b,t,d] = softplus(sum_{k} x[b,t,k] * W_delta[k,d] + b[d]). This is a matvec per t. Over L=2048 steps, that's 2048 * 256 = 524k MACs per (b,d). For 16 d-channels in a block: 8.4M MACs. With 256 threads, that's 32k MACs per thread. The threads could cooperate.
But this is getting complex. Let me just do the precomputed version first and see if it's fast enough.
Actually, let me reconsider the block configuration. The H100 has very high memory bandwidth, and the bottleneck for the scan is:
- Reading delta, x, B, C
- Computing exp
- Computing h updates
- Writing y
For 4096 (b,d) scans, each of length L=2048:
- Per step: 2 reads for delta, x (strided), 2 reads for B, C (2 * N contiguous floats), some compute
- Total reads: 4096 * 2048 * (2 + 2N) floats = 4096 * 2048 * 34 = ~285M floats = 1.14 GB
- Total writes: 4096 * 2048 * 1 = 8.4M floats = 33.5 MB
- At 3.35 TB/s, 1.14 GB reads = 0.34 ms
Hmm, that's a rough estimate. At 0.34ms, plus some compute overhead. The PyTorch reference with loop is probably 100x slower due to Python overhead and memory allocation.
So we should get a huge speedup. But to get maximum speedup, we need to optimize the kernel.
Let me think about the exact memory access pattern more carefully.
For each (b,d) processed by one thread:
```cpp
for (int t = 0; t < L; t++) {
float delta_t = delta[b * L * D + t * D + d];
float x_t = x[b * L * D + t * D + d];
float B_tn = B[b * L * N + t * N + n];
float C_tn = C[b * L * N + t * N + n];
...
}
```
This is the pattern for each of the 16 threads in a pair.
For threads in the same pair:
- delta_t and x_t are the same for all 16 threads (same b, d, t)
- B_tn and C_tn vary with n
For the 32 threads in a warp (2 pairs):
- Pair0 at (b, d): reads delta[b, t, d], x[b, t, d], B[b, t, 0..15], C[b, t, 0..15]
- Pair1 at (b, d+1) or (b+1, d'): reads delta[...], x[...], B[b, t, 0..15], C[b, t, 0..15]
Wait, if the pairs are adjacent d in the same batch, then pair1 reads delta[b, t, d+1], x[b, t, d+1]. These are adjacent in memory to pair0's reads. Since adjacent threads in a warp read adjacent memory locations, the delta and x reads are coalesced.
For B and C, pair0 reads B[b, t, 0..15] and pair1 reads B[b, t, 0..15] (same values if same batch). Wait, if pair0 and pair1 are different d in the same batch, they both read the same B[b, t, :] and C[b, t, :]!
So in the warp, threads 0..15 read B[b, t, 0..15] and threads 16..31 read B[b, t, 0..15] too! That's a broadcast / duplicate read. The L1 cache can handle this.
Similarly for C.
Great, so the memory pattern is:
- Delta: 2 coalesced reads (delta[b,t,d] and delta[b,t,d+1])
- X: 2 coalesced reads
- B: 1 cache line read, broadcast to all 32 threads? Well, 16 threads each reading 4 bytes, total 64 bytes = 1 cache line. The other 16 threads read the same 64 bytes. L1 cache handles this.
- C: same as B.
So the effective memory bandwidth is quite good.
What if pair0 and pair1 are in different batches? Then B reads are different and require 2 cache lines. But we can avoid this by having blocks process all 16 pairs from the same batch.
Yes, let's do that. Each block handles 16 contiguous d-channels from one batch. That means:
- blockIdx.x determines the batch and starting d
- For block i: b = i / (D / pairs_per_block), d_start = (i % (D / pairs_per_block)) * pairs_per_block
Actually, with 16 pairs per block and D=256, there are 256/16 = 16 d-groups per batch. With B=16, total blocks = 16 * 16 = 256 blocks.
Each block handles one batch b and d = d_start .. d_start + 15.
This is perfect! All pairs in a block share the same batch, so all B and C reads are the same.
Block assignment:
```cpp
int d_groups_per_batch = D / pairs_per_block; // 256/16 = 16
int b = blockIdx.x / d_groups_per_batch;
int d_start = (blockIdx.x % d_groups_per_batch) * pairs_per_batch;
int d = d_start + pair_idx_in_block;
```
Yes! This is much better.
Now, within a block, since all pairs read the same B[b, t, :] and C[b, t, :], we can load them into shared memory once per step.
Shared memory approach:
- `__shared__ float B_shared[N];` // 16 floats = 64 bytes
- `__shared__ float C_shared[N];` // 64 bytes
- At step t, have the first N threads (or pair 0's threads) load B and C into shared memory
- `__syncthreads()` if needed
- All threads read from shared memory
But with only 16 floats, a warp read already fits in one cache line. So explicit shared memory might not help much. The L1 cache (which is physically close to shared memory on H100) should already handle this.
However, to reduce cache pressure and ensure consistency, explicit shared memory is cheap and might help. But with 16 floats, `__syncthreads()` every iteration adds overhead.
Actually, we can have one thread in each warp load B and C. Wait, no, all threads in the warp read the same B[b,t,n] for their specific n. In a warp with 2 pairs (32 threads) all from the same batch:
- Thread i (0..15 pair0) reads B[b,t,i]
- Thread j (16..31 pair1) reads B[b,t,j-16]
Together they read B[b,t,0..15]. This is exactly 16 contiguous floats = 64 bytes = 1 cache line. Perfect.
And since all warps in the block process the same batch and read the same B[b,t,:], they'll all hit the same cache line. With 8 warps in a block, all reading the same cache line, L1 handles this with broadcasts.
So no need for shared memory for B and C. Good.
What about delta and x? For the 16 pairs in a block, they read delta[b, t, d_start..d_start+15]. That's 16 contiguous floats. If the 16 pairs are in 8 warps (2 pairs per warp, 32 threads), the threads in each warp read:
- Pair0 (d): delta[b,t,d]
- Pair1 (d+1): delta[b,t,d+1]
Together: 2 contiguous floats per warp? No, in a 32-thread warp, threads read:
- t0..t15 (pair0): delta[b,t,d0]
- t16..t31 (pair1): delta[b,t,d0+1]
Wait, the warp has 32 threads but only 16+16 are active for two pairs. Each pair reads 1 float per thread. For pair0, 16 threads each read 1 float, but they're all the same address! (Same b, t, d). No, they're the same delta_t for the same (b,t,d). So 16 threads in pair0 all read the same delta[b,t,d0]. That's a broadcast.
Similarly, pair1 reads delta[b,t,d0+1].
In the warp: threads 0..15 read address X, threads 16..31 read address X+sizeof(float). The hardware will do: 1 broadcast for X (to threads 0..15), and 1 load for X+4 (to threads 16..31). Actually, since threads 0..15 all read the same address, it's a broadcast (1 transaction). Threads 16..31 read the same address X+4, another broadcast (1 transaction). Total 2 transactions per warp.
For 8 warps reading delta: 8 * 2 = 16 transactions of 4 bytes each = 64 bytes total. But we need to read 16 floats = 64 bytes. So we're reading efficiently. Wait, but the addresses are delta[b,t,d0], delta[b,t,d0+1], ..., delta[b,t,d0+15] for the 16 pairs across 8 warps. Each warp reads 2 unique floats. Total: 8 warps * 2 floats = 16 floats.
In terms of cache lines: delta[b,t,d_start..d_start+15] fits in one 64-byte cache line (16 floats * 4 bytes = 64 bytes). So all these reads are from 1 cache line. L1 broadcasts the line to all warps. Very efficient!
Similarly for x.
Great. So the memory access is essentially optimal.
Now, what about the writes to y? Only the leader threads write. For 16 pairs, 16 leader threads write y[b,t,d_start..d_start+15]. These are 16 contiguous floats. If the 16 leaders are spread across 8 warps (2 per warp at lane 0 and lane 16), the writes are:
- Warp 0: thread 0 writes y[b,t,d0], thread 16 writes y[b,t,d0+1]
- Warp 1: thread 0 writes y[b,t,d0+2], thread 16 writes y[b,t,d0+3]
- etc.
This results in scattered writes within a cache line. In a warp, only 2 threads write, to adjacent addresses. The hardware will merge these. Across 8 warps, there are 16 writes to 16 contiguous floats. L1 handles write coalescing.
Actually, warp-level writes: in warp 0, 2 threads write to adjacent addresses. That's 1 transaction of 8 bytes. Across 8 warps: 8 transactions to the same cache line, total 64 bytes written. This is fine.
Alternatively, if we use `__syncthreads()` and have one thread write all 16 values from shared memory, we'd need to store the results in shared memory first. Let's not do that.
OK, I think the kernel design is solid. Let me write it out.
One more thing: how to compute the index mapping.
```cpp
__global__ void selective_scan_kernel(
const float* __restrict__ x,
const float* __restrict__ delta,
const float* __restrict__ A_log,
const float* __restrict__ B,
const float* __restrict__ C,
const float* __restrict__ D_param,
float* __restrict__ y,
int B, int L, int D, int N
) {
const int pairs_per_block = 16; // blockDim.x / N where N=16, blockDim.x=256
const int threads_per_pair = 16;
int tid = threadIdx.x;
int pair_in_block = tid / threads_per_pair; // 0..15
int n = tid % threads_per_pair; // 0..15
int d_groups = D / pairs_per_block;
int b = blockIdx.x / d_groups;
int d = (blockIdx.x % d_groups) * pairs_per_block + pair_in_block;
if (b >= B) return;
float A_val = -expf(A_log[d * N + n]);
float h = 0.0f;
size_t xd_offset = (size_t)b * L * D + d;
size_t bc_offset = (size_t)b * L * N;
size_t y_offset = (size_t)b * L * D + d;
for (int t = 0; t < L; t++) {
float delta_t = delta[xd_offset + t * D];
float x_t = x[xd_offset + t * D];
float B_tn = B[bc_offset + t * N + n];
float C_tn = C[bc_offset + t * N + n];
float dA = expf(delta_t * A_val);
float dB = delta_t * B_tn;
float x_db = x_t * dB;
h = dA * h + x_db;
float contrib = h * C_tn;
// Warp shuffle reduction: sum 16 values
// Thread layout: pairs_per_block=16, each pair has 16 threads
// tid 0..255: pair=tid/16, n=tid%16
// Warp 0: tid 0..31 -> pairs 0,1 -> n 0..15, 0..15
// We need to reduce within each 16-thread group.
// With __shfl_down_sync and width=16:
// Pair 0 (tid 0..15) are in sub-group 0 of warp 0.
// Pair 1 (tid 16..31) are in sub-group 1 of warp 0.
// ...
contrib += __shfl_down_sync(0xFFFFFFFF, contrib, 8, 16);
contrib += __shfl_down_sync(0xFFFFFFFF, contrib, 4, 16);
contrib += __shfl_down_sync(0xFFFFFFFF, contrib, 2, 16);
contrib += __shfl_down_sync(0xFFFFFFFF, contrib, 1, 16);
if (n == 0) {
y[y_offset + t * D] = contrib + x_t * D_param[d];
}
}
}
```
Wait, will this compile and work correctly? Let me verify __shfl_down_sync.
`__shfl_down_sync(unsigned mask, T var, unsigned int delta, int width=warpSize)`
`mask`: must include all threads participating. `0xFFFFFFFF` means all threads in warp.
`delta`: offset within the sub-group.
`width`: partition the warp into sub-groups of this width. If width=16, warp is split into two 16-thread groups. The `delta` is applied within each sub-group.
For pair 0 (tid 0..15): `__shfl_down_sync(0xFFFFFFFF, contrib, 8, 16)` gets contrib from thread within sub-group at offset +8. So thread 0 gets from thread 8, thread 1 from thread 9, ..., thread 7 from thread 15. Threads 8..15 get undefined? No, `__shfl_down_sync` has a sync requirement but width constrains the exchange. Actually, for threads 8..15, the source would be thread 16..23, but those are in a different sub-group. The result might be the value from thread X+delta within the same sub-group, but if X+delta >= width, what happens?
The documentation says: "__shfl_down_sync() calculates a source lane ID by adding delta to the caller's lane ID: srcLane = laneId + delta. If srcLane is out of range [0:width) or srcLane is not in mask, the value returned is the calling thread's own var."
Wait, that's for when width is specified. Let me check: "If width is less than warpSize then each subsection of the warp behaves as a separate entity with a starting logical lane ID of 0."
So for width=16:
- Sub-group 0: logical lanes 0..15 (physical lanes 0..15)
- Sub-group 1: logical lanes 0..15 (physical lanes 16..31)
Thread 8 (physical) is logical lane 8 in sub-group 0. `__shfl_down_sync(..., 8, 16)` -> srcLane = 8 + 8 = 16. But 16 is out of range [0, 16), so it returns its own value. That's correct - in a butterfly reduction, we want threads 0-7 to get from 8-15, and threads 8-15 to keep their own.
So:
```cpp
contrib += __shfl_down_sync(0xFFFFFFFF, contrib, 8, 16); // 0-7 += 8-15, 8-15 unchanged
contrib += __shfl_down_sync(0xFFFFFFFF, contrib, 4, 16); // 0-3 += 4-7, etc.
contrib += __shfl_down_sync(0xFFFFFFFF, contrib, 2, 16);
contrib += __shfl_down_sync(0xFFFFFFFF, contrib, 1, 16);
```
After this, logical lane 0 of each sub-group has the sum. Physical lanes 0 and 16 have the sums.
Then `if (n == 0)`: for pair 0, n=0 maps to thread 0 (physical lane 0). For pair 1, n=0 maps to thread 16 (physical lane 16). Both correctly are the leaders of their sub-groups. Perfect.
This looks correct.
One more optimization: I can prefetch delta_t+1, x_t+1, B_t+1, C_t+1 while computing step t. This is a classic software pipelining technique. But with registers already being used, and the compiler usually handling this, let me not do it manually unless needed.
Another idea: since N=16, we can unroll the N=16 dimension. But we're already using threads for N parallelism. Unrolling doesn't help here.
What about using half2 or bf16? The problem says fp32, so stick to float.
Wait, let me check: "Produce numerically equivalent outputs at fp32 precision." This means the outputs should match in fp32, but the internal compute can be in higher precision? Actually, "at fp32 precision" usually means the output is fp32 and we match within fp32 tolerance. expf is accurate to ~1 ULP. The reference uses torch.exp which is also accurate. Should be fine.
One concern: A_log might cause some numerical issues. The reference computes `A = -torch.exp(self.A_log)` once, then uses it in the loop. Our kernel computes A_val once per thread. Same thing, just distributed.
Now, let me also think about whether we should vectorize the loop over t. Since L=2048, and the kernel doesn't know L at compile time (it's passed as argument), unrolling is limited. We could process 2 or 4 t-steps per loop iteration:
```cpp
#pragma unroll 4
for (int t = 0; t < L; t++) {
...
}
```
`#pragma unroll 4` should be safe and beneficial.
Actually, the compiler does this automatically if it can. But we can add it.
Also, `A_log` is read from global memory. Since A_log[d*N + n] is only read once per thread and then kept in a register, it's fine.
Let me think about whether to use `float` or `__half`. The specs say fp32 output, so float is fine.
Now let me write the complete ModelNew code.
Actually, one more thing: the kernel function signature. We need to expose it to pybind11 via `load_inline`. The function should be a host function that takes tensors and calls the kernel.
```cpp
void selective_scan_forward(
at::Tensor x,
at::Tensor delta,
at::Tensor A_log,
at::Tensor B,
at::Tensor C,
at::Tensor D_param,
at::Tensor y,
int B, int L, int D, int N
) {
const int block_size = 256;
const int pairs_per_block = block_size / N; // 16
int d_groups = D / pairs_per_block;
int num_blocks = B * d_groups;
selective_scan_kernel<<<num_blocks, block_size>>>(
x.data_ptr<float>(),
delta.data_ptr<float>(),
A_log.data_ptr<float>(),
B.data_ptr<float>(),
C.data_ptr<float>(),
D_param.data_ptr<float>(),
y.data_ptr<float>(),
B, L, D, N
);
}
```
Wait, we should also check that D is divisible by pairs_per_block (=16). D=256, so yes.
And we need to handle non-contiguous tensors. For simplicity, assume contiguous. We can add `.contiguous()` in the Python wrapper.
Also, need to include appropriate headers.
One more check: is D_param [D]? Yes.
Now, in ModelNew.__init__, we initialize from the same arguments.
In ModelNew.forward, we:
1. Compute delta = F.softplus(F.linear(x, self.proj_delta.weight, self.proj_delta.bias))
2. Compute B_t = F.linear(x, self.proj_B.weight)
3. Compute C_t = F.linear(x, self.proj_C.weight)
4. Allocate y = torch.empty_like(x)
5. Call kernel
6. Return y
Wait, the reference uses `self.proj_delta(x)` which does the linear + bias. We can just call it.
But actually, a small optimization: instead of using F.linear, we can just use the nn.Linear layers.
```python
delta = torch.nn.functional.softplus(self.proj_delta(x))
B_t = self.proj_B(x)
C_t = self.proj_C(x)
```
Now for the kernel code, I need to be careful about the C++ function signature for `load_inline`. The function should take `at::Tensor` arguments.
Let me draft the full code:
```python
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.cpp_extension import load_inline
cuda_source = """
#include <torch/extension.h>
#include <cuda_runtime.h>
__global__ void selective_scan_kernel(
const float* __restrict__ x,
const float* __restrict__ delta,
const float* __restrict__ A_log,
const float* __restrict__ B,
const float* __restrict__ C,
const float* __restrict__ D_param,
float* __restrict__ y,
int B, int L, int D, int N
) {
const int threads_per_pair = 16;
const int pairs_per_block = blockDim.x / threads_per_pair;
int tid = threadIdx.x;
int pair_in_block = tid / threads_per_pair;
int n = tid % threads_per_pair;
int d_groups = D / pairs_per_block;
int b = blockIdx.x / d_groups;
int d = (blockIdx.x % d_groups) * pairs_per_block + pair_in_block;
if (b >= B) return;
float A_val = -expf(A_log[d * N + n]);
float h = 0.0f;
size_t xd_stride_D = (size_t)D;
size_t bc_stride_N = (size_t)N;
size_t xd_batch_offset = (size_t)b * L * D;
size_t bc_batch_offset = (size_t)b * L * N;
size_t x_base = xd_batch_offset + d;
size_t b_base = bc_batch_offset;
size_t y_base = xd_batch_offset + d;
for (int t = 0; t < L; t++) {
size_t t_xd = x_base + (size_t)t * D;
size_t t_bc = b_base + (size_t)t * N;
float delta_t = delta[t_xd];
float x_t = x[t_xd];
float B_tn = B[t_bc + n];
float C_tn = C[t_bc + n];
float dA = expf(delta_t * A_val);
float dB = delta_t * B_tn;
float x_db = x_t * dB;
h = dA * h + x_db;
float contrib = h * C_tn;
contrib += __shfl_down_sync(0xFFFFFFFF, contrib, 8, 16);
contrib += __shfl_down_sync(0xFFFFFFFF, contrib, 4, 16);
contrib += __shfl_down_sync(0xFFFFFFFF, contrib, 2, 16);
contrib += __shfl_down_sync(0xFFFFFFFF, contrib, 1, 16);
if (n == 0) {
y[y_base + (size_t)t * D] = contrib + x_t * D_param[d];
}
}
}
void selective_scan_forward(
at::Tensor x,
at::Tensor delta,
at::Tensor A_log,
at::Tensor B,
at::Tensor C,
at::Tensor D_param,
at::Tensor y,
int B,
int L,
int D,
int N
) {
const int threads_per_pair = 16;
const int block_size = 256;
const int pairs_per_block = block_size / threads_per_pair;
int d_groups = D / pairs_per_block;
int num_blocks = B * d_groups;
selective_scan_kernel<<<num_blocks, block_size>>>(
x.data_ptr<float>(),
delta.data_ptr<float>(),
A_log.data_ptr<float>(),
B.data_ptr<float>(),
C.data_ptr<float>(),
D_param.data_ptr<float>(),
y.data_ptr<float>(),
B, L, D, N
);
}
"""
```
Hmm wait, I need to check: does `__shfl_down_sync` with width=16 and mask=0xFFFFFFFF work correctly when the calling thread is not in the mask? All threads in the warp are in the mask (0xFFFFFFFF = all 32), and width=16 splits into two sub-groups. The `__shfl_down_sync` docs say: "mask must include all threads participating in the call, or the behavior is undefined." Since all threads participate, 0xFFFFFFFF is correct.
One potential issue: I need `n == 0` to identify the leader. For pair_in_block=0, threadIdx.x=0 has n=0, so it's a leader. For pair_in_block=1, threadIdx.x=16 has n=0, also leader. Good.
Let me also consider if D=256 is not divisible by pairs_per_block=16. 256/16 = 16 exactly. Good.
Now, the Python ModelNew:
```python
class ModelNew(nn.Module):
def __init__(self, d_model: int, d_state: int):
super(ModelNew, self).__init__()
self.d_model = d_model
self.d_state = d_state
self.A_log = nn.Parameter(torch.randn(d_model, d_state))
self.proj_B = nn.Linear(d_model, d_state, bias=False)
self.proj_C = nn.Linear(d_model, d_state, bias=False)
self.proj_delta = nn.Linear(d_model, d_model, bias=True)
self.D = nn.Parameter(torch.ones(d_model))
def forward(self, x: torch.Tensor) -> torch.Tensor:
B, L, D = x.shape
N = self.d_state
delta = F.softplus(self.proj_delta(x))
B_t = self.proj_B(x)
C_t = self.proj_C(x)
y = torch.empty_like(x)
selective_scan_forward(
x.contiguous(),
delta.contiguous(),
self.A_log.contiguous(),
B_t.contiguous(),
C_t.contiguous(),
self.D.contiguous(),
y,
B, L, D, N
)
return y
```
Wait, I should make sure the tensors are contiguous. The `contiguous()` call ensures this. But it might make a copy if the tensor is already contiguous. `x` should be contiguous from `torch.randn`, but `.contiguous()` is cheap.
Actually, a better approach: use `x = x.contiguous()` or just ensure inputs to the kernel are contiguous.
One optimization: instead of passing `x` as a separate tensor and reading it, could we avoid it? We need x for x_db and for the skip connection. So x needs to be read. It's already in global memory.
Now, a concern about numerical equivalence: the PyTorch reference uses double precision for intermediate computations? No, by default PyTorch uses float32 and all operations are float32. Our kernel uses float32 explicitly. Should match.
Actually, `torch.exp` in PyTorch might use a different exp implementation than CUDA `expf`. But both should be accurate to within ~1 ULP. The tolerance is 1e-4 for fp32, which is generous compared to machine epsilon of ~1e-7. So matching exp exactly doesn't matter.
One thing: the reference `dA` is computed for all t at once using `torch.exp(delta.unsqueeze(-1) * A.unsqueeze(0).unsqueeze(0))`. This computes `delta * A` first, then applies `exp` elementwise. In our kernel, we do `expf(delta_t * A_val)`. Same mathematical operation, same order.
However, `torch.exp` for tensors might use vectorized implementations that have slightly different rounding than scalar `expf`. But again, tolerance 1e-4 should be fine.
Wait, let me also think about the dimensions more carefully.
The reference:
```
dA = torch.exp(delta.unsqueeze(-1) * A.unsqueeze(0).unsqueeze(0))
# delta: [B, L, D, 1]
# A: [1, 1, D, N]
# dA: [B, L, D, N]
```
In kernel: `expf(delta_t * A_val)` where `A_val = -expf(A_log[d * N + n])`. A in reference is `-exp(A_log)`. So `delta_t * A_val = delta_t * (-exp(A_log[d,n]))`. Matches.
Then `dB = delta.unsqueeze(-1) * B_t.unsqueeze(2)`
# delta: [B, L, D, 1]
# B_t: [B, L, 1, N]
# dB: [B, L, D, N]
In kernel: `dB = delta_t * B_tn`. Matches.
`x_db = x.unsqueeze(-1) * dB`
# x: [B, L, D, 1]
# dB: [B, L, D, N]
# x_db: [B, L, D, N]
In kernel: `x_db = x_t * dB`. Matches.
`h = dA[:, t] * h + x_db[:, t]`
In kernel: `h = dA * h + x_db`. Matches.
`y_t = (h * C_t[:, t].unsqueeze(1)).sum(dim=-1)`
In kernel: `contrib = h * C_tn`, then warp sum. Matches.
`y = y + x * self.D`
In kernel: `contrib + x_t * D_param[d]`. Matches.
Great.
Now, one more optimization concern: the kernel has a loop of L with `__shfl_down_sync` in every iteration. Shuffle instructions are fast but not free. L=2048 means 2048 shuffle sequences per thread. That's fine.
Can we reduce the number of shuffles? What if we accumulate h and only do the dot product at the end of every K steps? No, because y needs to be output at every t.
We can try to use warp-level matrix operations or tensor cores for the N=16 reduction. But with 16 elements, warp shuffle is optimal. Tensor cores have minimum dimensions.
One thing I want to check: using TMA (Tensor Memory Accelerator) on H100. But TMA is mainly for global-to-shared transfers in a structured way. Our memory access pattern is simple enough that regular loads are fine. TMA would add complexity for minimal gain here.
How about register pressure? Each thread needs:
- A_val (1)
- h (1)
- delta_t (1)
- x_t (1)
- B_tn (1)
- C_tn (1)
- dA, dB, x_db, contrib (4)
- Loop counter, indices (~8)
Total ~20 registers. For 256 threads per block, that's 5120 registers per block. Each SM has 64K registers, so can fit ~12 blocks per SM concurrently. With 256 blocks total and 132 SMs, that's ~2 blocks per SM. So ~24 warps per SM. Good occupancy.
On H100, the warp scheduler can handle many warps. With 24 warps per SM, that's decent.
Wait, can we increase occupancy by reducing block size? If we use 128 threads (4 warps) per block:
- 4096 pairs / 8 pairs per block = 512 blocks.
- 512 blocks / 132 SMs = ~4 blocks per SM.
- Registers per block: 128 * 20 = 2560.
- Blocks per SM: 64K / 2560 = ~25. But limited by max blocks per SM = 32.
- So 4 blocks per SM concurrently, 16 warps per SM.
Compared to 256 threads:
- 256 blocks / 132 = ~2 blocks per SM.
- 16 warps per SM.
Same warp count per SM! But 512 blocks might have better load balancing and less tail latency. Let me switch to 128 threads per block, which gives 512 blocks.
Actually wait, 128 threads with N=16 means 8 pairs per block. 4096 pairs / 8 = 512 blocks. 512 blocks / 132 SMs = ~3.9 blocks per SM. On H100, max 32 blocks per SM. So we're well under.
But each block now has 4 warps instead of 8. With 4 blocks per SM, that's 16 warps per SM. Same as before.
Actually, let me compute more carefully. If a block has 8 warps (256 threads), and 2 blocks per SM, that's 16 warps. If a block has 4 warps (128 threads), and 4 blocks per SM, that's 16 warps. The difference is:
- More blocks = better load balancing across SMs, especially if some blocks finish earlier.
- Fewer threads per block = less shared memory per block (not relevant here, we don't use much).
- The main factor for latency hiding is the number of warps per SM, which is similar.
I'll go with 256 threads per block for simplicity, but 128 is also fine. Actually, let me use 128 threads to get 512 blocks, which is closer to the "magic number" for occupancy. But 256 blocks is also fine - GPUs often perform well with hundreds of blocks.
Actually, on H100, we want at least as many blocks as SMs (132) to utilize all SMs. 256 > 132, so we're good. 512 would give even better distribution.
Hmm, another thing: with 256 threads per block and 256 blocks, the total threads is 65536. The GPU can handle many more. With 128 threads and 512 blocks, also 65536 threads.
Let me try 128 threads first. It gives 512 blocks which might have slightly better SMs utilization.
With block_size=128, pairs_per_block = 128/16 = 8. d_groups = 256/8 = 32. num_blocks = 16 * 32 = 512.
In the kernel:
```cpp
const int threads_per_pair = 16;
const int pairs_per_block = 8; // blockDim.x / threads_per_pair
```
Wait, but then in a warp of 32 threads, we'd have pairs 0 and 1 (16 threads each). But there are 8 pairs per block = 4 warps. That's fine.
Let me stick with 256 threads (16 pairs, 8 warps) for now. It doesn't matter too much. Actually 128 is also fine.
Let me also consider processing more than one d-channel per thread. If a single thread processes N=16 elements sequentially, then a block of 128 threads handles 128 (b,d) pairs. 4096/128 = 32 blocks. That's way too few.
Alternatively, each thread handles M n-elements, so each pair has 16/M threads. For M=2, each pair has 8 threads, block of 256 handles 32 pairs. 4096/32 = 128 blocks. Better, but the shuffle reduction becomes more complex.
Actually, we can have threads in a warp handle multiple (b,d) pairs if we reduce the number of threads per pair. For example, if we use 8 threads per pair (each thread handles 2 n values), then a warp handles 4 pairs. But then each thread needs 2 h values, and the shuffle reduction needs to happen across 8 threads with each thread having 2 values.
This is possible but more complex. Let me not complicate things. The current design with 16 threads per pair and 256 threads per block should work well.
Actually, I realize there might be an issue with `__shfl_down_sync`. If the kernel is compiled for compute capability that doesn't support it... but H100 is sm_90, which definitely supports it. And `load_inline` will use the appropriate architecture.
Let me check: `__shfl_down_sync` requires CUDA 9.0+ and compute capability 3.5+ for `_sync` variants? Actually, `__shfl_down_sync` requires compute capability 3.5+ for the sync variant and is available in all modern GPUs. H100 (sm_90) definitely has it.
One more thing: I need to make sure I `#include <cuda_runtime.h>` and use proper ATen headers. `torch/extension.h` should pull in what's needed.
Also, `load_inline` with `cpp_sources` containing CUDA code. CUDA code in `cpp_sources` works because NVCC handles `.cu` files. The extension system compiles it.
Let me also think about whether I should process the linear projections inside the kernel to save memory bandwidth.
Consider: the kernel currently reads x, delta, B, C, A_log, and D. Then writes y. But delta, B, C are computed from x. If we fused, we'd only read x, A_log, D and write y. And we'd read the 3 weight matrices.
Delta projection weights: [D, D] = 256*256 = 65k floats = 256KB
B projection weights: [D, N] = 256*16 = 4k floats = 16KB
C projection weights: same.
Total weight reads: ~288KB, read once per block? No, each block needs a subset of weights. If each block handles 16 d-channels, it needs all of B and C weights (since they don't depend on d), and 16 columns of delta weights plus 16 bias values.
If all 512 blocks read the B weights (16KB each), that's 8MB of B weight reads total. And 8MB of C weight reads. This would be in addition to reading x. That's actually more memory traffic than just reading precomputed B and C (which are 32KB per block, 16MB total for all blocks... wait, 512 blocks * 32KB = 16MB. But B and C are the same tensor, read from L2 cache. Total actual memory reads: B is 2MB, C is 2MB, read multiple times by different blocks but cached in L2. So actual DRAM traffic is just ~4MB for B and C. For weights, also ~4MB total. So they're similar.
For x: reading x is 16 * 2048 * 256 = 8MB. If we precompute delta, B, C, we also read them: delta is 8MB, B is 2MB, C is 2MB. Total additional reads: 12MB. If we fuse, we read x once (8MB) and weights once (~0.3MB, mostly cached). That's a significant saving: 8MB vs 20MB of reads!
And for the scan itself, we still need to read delta and B, C in the loop. Wait, if we fuse, we compute delta, B, C on the fly in the loop. In each step t, for each (b,d):
- Read x[b,t,d]
- Compute delta from x[b,t,:] (matvec of length D)
- Compute B from x[b,t,:] (matvec of length D, output N)
- Compute C from x[b,t,:] (matvec of length D, output N)
- Then do the scan step
In the precomputed version, per step t per (b,d):
- Read delta[b,t,d], x[b,t,d], B[b,t,:], C[b,t,:]
- Then scan step
The precomputed version reads:
- delta_t: 1 float
- x_t: 1 float
- B_t: N floats (but shared across d in same batch)
- C_t: N floats (shared across d in same batch)
Per step per pair: 2 floats + 2N floats amortized across d in same batch.
In the fused version, per step per pair:
- x_t: 1 float
- Need x[b,t,0..D-1] to compute delta_t, B_t, C_t. But this is shared across all d in the same batch.
If a block handles multiple d for the same b, and all threads need x[b,t,0..D-1], they can load it cooperatively into shared memory. D=256 floats = 1KB. That's tiny! With 8 or 16 pairs per block, we just need 1KB of shared memory for the entire x row per batch.
Then all threads in the block read x[b,t,0..D-1] from shared memory to compute delta_t, B_t, C_t for their specific d and n.
Wait, this is getting complex but potentially much faster because:
1. Global memory reads: only read x[b,t,:] once per step per block
2. No reads of delta, B, C from global memory
3. Weight matrices are small and stay in L1/L2 cache
For each step t, the block:
1. Loads x[b, t, d_start..d_start+pairs_per_block-1] into registers (already being read by each thread for its own d)
2. But to compute delta for all d in the block, each thread needs all of x[b,t,:]! Because delta_t = sum_k x[b,t,k] * W_delta[k,d] for each d.
Wait, in the original model, `self.proj_delta` is `nn.Linear(d_model, d_model)`. So delta[b,t,d] = sum_k x[b,t,k] * W_delta[k,d] + b[d].
For a block with 16 d-channels, the 16 threads in pair i need x[b,t,:] to compute delta[b,t,d_i]. Specifically, each of the 16 threads in a pair reads delta[b,t,d] which was computed from the full x[b,t,:]. But those 16 threads all share the same delta[b,t,d] (since they have the same d).
This is really a matvec: x_row @ W_delta_cols. With 256 threads in a block handling 16 d-channels (16 pairs * 16 threads each), only 16 unique delta values are needed. These could be computed by 16 threads, or cooperatively.
Actually, the simplest approach for fusing: have the block compute delta, B, C for the current t in shared memory using cooperative methods, then do the scan step.
For B and C (which are [N] vectors, same for all d in the batch):
- B[b,t,n] = sum_k x[b,t,k] * W_B[k,n]
- C[b,t,n] = sum_k x[b,t,k] * W_C[k,n]
These are matvecs of length 256 to 16 outputs. A block of 256 threads can compute this in one step:
- Thread i computes partial products for x_i * W_B[i, n_i] and accumulates. Actually, for a matvec, each thread can handle some output elements.
For B: 16 outputs, 256 threads. We can have 16 threads handle each output, doing 256/16 = 16 MACs per thread, then warp-shuffle reduce. Or use 256 threads to compute all 16 outputs cooperatively.
This is getting complex. Let me see if the simple version (precomputed delta, B, C) is fast enough. The precomputed version reads about 20MB from memory and the scan has light compute. The memory bandwidth is 3.35 TB/s, so 20MB takes ~6 microseconds. The compute overhead for 134M thread-steps is tiny. So total kernel time might be ~10-20 microseconds.
The PyTorch reference allocates 500MB+ tensors and does Python loops. It probably takes milliseconds. So even the simple kernel should give 100x+ speedup.
But to be safe and get maximum speedup, I should try the simplest correct kernel first, validate correctness, then optimize.
Actually, I'm worried about one thing: the reference does `dA = torch.exp(delta.unsqueeze(-1) * A.unsqueeze(0).unsqueeze(0))` which computes a [B,L,D,N] tensor. That's 134M elements. Computing and storing this in PyTorch is very slow. Our kernel avoids this entirely. Even with precomputed delta, B, C, we avoid the huge intermediate.
So let me just go with the precomputed version.
Wait, one more optimization: can I avoid storing delta, B, C as separate tensors? Instead of calling softplus and the linear layers and then passing to a kernel, I could compute them in PyTorch but use the kernel to consume them directly. But PyTorch will still allocate and write them. There's no way around that without fusing.
Actually, PyTorch might materialize these anyway. And the scan kernel reading them is limited by bandwidth.
Let me test with a basic kernel first.
One thing I want to check: should I use float4 for the B and C loads? Not straightforward since each thread needs one element based on `n`.
Wait, 16 threads in a pair read B[b,t,0..15] contiguously. They could each read a float4 if N were larger. But for a single float per thread, there's no vectorization.
But there is an alternative: process N using vector types. Since N=16, a single thread could load float4 B_t[n:n+4] and process 4 n-values. That would mean fewer threads per pair.
For instance, each thread handles 4 n-values. Then a pair needs 4 threads instead of 16. The block of 256 threads could handle 64 (b,d) pairs. 4096/64 = 64 blocks. Too few blocks.
Or a thread handles 2 n-values: pair needs 8 threads, block handles 32 pairs, 4096/32 = 128 blocks. Better.
With 4 threads per pair and each thread handling N/4 = 4 elements:
- Register usage per thread: h[4], A[4], B[4], C[4], etc. ~20 registers still.
- Block of 256 threads: 64 pairs.
- Blocks: 4096/64 = 64 blocks. Still low.
- But each thread does 4x more compute per step and 4 loads. The shuffle reduction is across 4 threads, much fewer shuffles!
Wait, for sum across 4 threads in a warp: `__shfl_down_sync(mask, val, 2, 4)` and `__shfl_down_sync(mask, val, 1, 4)`. Only 2 shuffles instead of 4. But with 64 blocks, we have underutilization.
Actually, with 64 blocks of 256 threads, we have 64 warps per block? No, 256 threads = 8 warps. Total warps = 64 * 8 = 512 warps. On 132 SMs, that's ~4 warps per SM. That's quite low.
If instead we use 8 threads per pair (each handling 2 n-values):
- Block of 256 threads: 32 pairs.
- Blocks: 4096/32 = 128 blocks.
- Warps per block: 8. Total warps: 128 * 8 = 1024. ~8 warps per SM. Much better.
With 16 threads per pair (each handling 1 n-value):
- Block of 256 threads: 16 pairs.
- Blocks: 256 blocks.
- Total warps: 256 * 8 = 2048. ~15 warps per SM. Good.
So 16 threads per pair is actually better for parallelism. The tradeoff is more shuffle instructions but better SM utilization.
What about vectorization with 16 threads per pair? We can't really vectorize the loads since each thread gets a unique n.
One more idea: use 2 warps per pair. 32 threads per (b,d), each handling 1 n-value (but half idle) or each handling multiple elements. But 32 threads per pair halves the number of pairs. Not good.
I'll stick with 16 threads per pair.
Let me also think about block size 128:
- 8 pairs per block.
- Blocks: 512.
- Warps: 512 * 4 = 2048. ~15 warps per SM. Same as 256 threads, 256 blocks.
- But 512 blocks might have better scheduling.
I'll use 128 threads per block to get 512 blocks.
Actually wait, with block size 128 and 8 warps per block (no, 128 threads = 4 warps). 512 blocks * 4 warps = 2048 warps. 2048 / 132 = ~15.5 warps per SM.
With block size 256 and 16 pairs: 256 blocks * 8 warps = 2048 warps. Same total warps, same warps per SM.
The difference is blocks per SM: 512/132 ≈ 3.9 vs 256/132 ≈ 1.9. More blocks per SM is generally better for hiding latency and keeping the scheduler busy.
So 128 threads / 512 blocks is slightly better.
Let me also consider block size 64:
- 4 pairs per block.
- Blocks: 4096/4 = 1024 blocks.
- Warps: 1024 * 2 = 2048 warps.
- Blocks per SM: 1024/132 ≈ 7.8.
Even better distribution! But 64 threads per block is on the smaller side. Some GPUs have overhead for small blocks. But on H100, 64-thread blocks are fine.
However, with 2 warps per block and 1024 blocks, the block launch overhead is higher. But 1024 blocks is well within limits.
The main concern is whether a small block can saturate a warp scheduler. Each SM has multiple warp schedulers (4 on H100?). With 2 warps per block and 8 blocks per SM, that's 16 warps per SM. That's plenty for the schedulers.
I'll go with 128 threads per block (4 warps, 8 pairs, 512 blocks). It's a good balance.
Actually, let me compute optimal block size more carefully. The kernel has:
- 1 loop of L=2048 iterations
- Each iteration has: 4 loads, 1 expf, ~6 FMAs, 4 shuffles, 1 store (only for n==0)
The loop has a data dependency on h. So ILP is limited. The memory loads have some latency. With more warps per SM, we can hide latency better.
At 128 threads (4 warps) per block and 512 blocks, on 132 SMs: ~4 blocks per SM = 16 warps per SM. H100 can run up to 64 warps per SM (64KB registers / (256 regs/warp) = 256 warps... wait no, 2048 threads / 32 = 64 warps per SM max). So 16 warps is well under the limit. Can we get more?
With 64 threads (2 warps) per block and 1024 blocks: ~8 blocks per SM = 16 warps per SM. Same total warps per SM.
With 256 threads (8 warps) per block and 256 blocks: ~2 blocks per SM = 16 warps per SM. Same again.
So for any block size that gives 8 or 16 pairs per block, we get ~16 warps per SM, which is about 25% of max occupancy. Can we do better?
To get more warps per SM, we'd need more total warps. But total warps is total threads / 32 = (B * D * N / N) / 32 = B*D/32 = 4096/32 = 128 warps. Wait, that's not right.
Total thread instances: we have one thread per (b,d,n) = 65536 threads = 2048 warps. But these warps are grouped into blocks. The total number of warps is fixed at 2048. On 132 SMs, that's 2048/132 ≈ 15.5 warps per SM.
Oh! I see. If we use one thread per (b,d,n), total threads = 65536, total warps = 2048. The block size only affects how these 2048 warps are grouped. The warps per SM is always ~15.5 regardless of block size (assuming equal distribution).
But wait, for 128-thread blocks (4 warps), 512 blocks: 512 * 4 = 2048 total warps. Correct.
For 256-thread blocks (8 warps), 256 blocks: 256 * 8 = 2048 warps. Correct.
So the total parallelism is fixed. The difference is just scheduling granularity.
15.5 warps per SM out of 64 max is about 24% theoretical occupancy. Is that enough? Probably yes, because the kernel is memory-bound and each warp has a lot of work. But could we get higher occupancy? Only if we use more threads, and we can't because N=16 is the natural parallelism unit per (b,d).
Unless... we process each n-value with multiple threads (redundant work) or use a thread block to process one (b,d) with all N handled redundantly. But that reduces total parallel instances.
For example, if we use a warp of 32 threads for each (b,d), with 16 active threads each handling 1 n-value and 16 idle, total warps = 4096 (since there are 4096 pairs and 1 warp per pair). Total threads = 4096 * 32 = 131072. Total warps = 4096. Warps per SM = 4096/132 ≈ 31. That's double!
But half the threads in each warp are idle. The scheduler issues instructions to the whole warp, and the inactive threads are masked out. So effective work per instruction is halved. But total memory bandwidth might increase because there are more warps to overlap.
Actually, the active threads in a warp are contiguous: lanes 0..15 for pair0, lanes 16..31 for pair1 (if we process 2 pairs per warp). Wait, if I process only 1 pair per warp, 16 threads active and 16 idle in each warp. Then total warps = 4096.
With 2 pairs per warp (32 active threads), total warps = 2048.
So 1 pair per warp gives 2x the warp count but 2x more idle threads. The effective work per SM is similar. However, having 2x warps means the warp scheduler has more warps to switch between, potentially better latency hiding for memory operations.
But the actual number of memory requests is the same (same work, just more masked threads). And the SM has limited resources. More warps means more register pressure per SM? No, total registers for all warps on an SM can't exceed 64K.
With 1 pair per warp: each warp has 32 threads but only 16 do work. The compiler might still allocate 32 thread slots. Register count per SM is per thread, not per active thread. So 4096 warps * 32 threads * ~20 registers = 2.6M registers needed across all SMs. But each SM only has 64K registers, so it can hold at most 64K / (32 * 20) = 100 threads = ~3 warps concurrently. That's a huge occupancy problem!
Actually, registers are allocated per thread, not per warp. If a kernel uses R registers per thread, each block of T threads uses T*R registers. An SM can run at most floor(64K / (T*R)) blocks concurrently.
For our kernel with 20 registers per thread:
- 128 threads per block: 2560 registers per block. Max blocks per SM = 64K/2560 = 25. But capped at 32.
- 256 threads per block: 5120 registers per block. Max blocks per SM = 12.
Total blocks = 512 (128-thread) or 256 (256-thread). With 132 SMs:
- 128-thread: 512/132 ≈ 4 blocks per SM. Well under the max of 25. So actual concurrent blocks = 4. Warps per SM = 4*4 = 16.
- 256-thread: 256/132 ≈ 2 blocks per SM. Well under max of 12. Actual warps = 2*8 = 16.
Same again.
What if we reduce register count to 10? Then 128-thread blocks use 1280 regs, max 50 blocks per SM. But we only have 4 blocks per SM. No improvement.
Conclusion: the parallelism is fundamentally limited by the 65536 total (b,d,n) elements and B*D=4096 (b,d) instances. With ~132 SMs, we get ~500 threads per SM or ~15 warps. This is decent for a memory-bound kernel.
Let me just go with 256 threads per block for simplicity.
One more thing: should I consider splitting the L dimension and having multiple thread blocks per (b,d)? Probably not worth it for L=2048. The sequential loop per thread is fine.
Wait, I just realized something important. We can parallelize over L by using a parallel scan algorithm like the Blelloch scan. But Blelloch scan is for associative operations, and we have an affine recurrence h_t = a_t * h_{t-1} + b_t.
The scan composition is: (a2, b2) ∘ (a1, b1) = (a2*a1, a2*b1 + b2).
If we define each step as a pair (a_t, b_t), then the prefix scan of these pairs gives us the composed transformation from step 0 to t. We can apply this to initial state h_0 = 0.
This is associative and can be parallelized! But implementing aBlelloch scan for 65536 instances each of length 2048 with N=16 is complex. Each (b,d,n) is an independent scan instance. We'd need to do a segmented/parallel scan over L for each of the 65536 instances.
With 65536 independent scans, and L=2048, a naive parallel scan would use ~L/2 log L = ~11k threads per scan instance, which is overkill. More practically, we could have each block handle one (b,d) and cooperatively scan L using shared memory between the 16 threads in the block. But with 4096 (b,d) instances and only 16 threads per instance, the block size would be 16 threads, giving 4096 blocks.
For a block of 16 threads scanning L=2048:
- Load all a_t, b_t for t=0..2047 into shared memory? That's 2048 * 2 * 16 * 4 bytes ≈ 262KB per block. Too much shared memory (max 227KB).
- So we need to do it in chunks.
A more practical approach: use Kogge-Stone or Brent-Kung scan in shared memory for chunks of say 512 elements, then combine chunks. This requires multiple passes and is complex.
For L=2048, the sequential per-thread scan is probably faster than the overhead of a parallel scan with small blocks. Parallel scan shines for very long sequences. 2048 is on the border, but with 13.9M elements total to scan and memory-bound behavior, the sequential approach is likely optimal.
Actually, there's a hybrid approach: process L in chunks of say K steps per thread, where K = L / threads_per_block. But with only 16 threads per (b,d), each thread would handle 128 steps. We could use shared memory to store intermediate states and do a tree-based reduction. But again, complex.
I'll stick with the single-thread sequential scan.
Let me write the code and test it.
Actually, before writing the full code, let me think about one more optimization: `__restrict__` and pointer arithmetic to avoid repeated indexing.
```cpp
const float* __restrict__ x_ptr = x + b * L * D + d;
const float* __restrict__ delta_ptr = delta + b * L * D + d;
const float* __restrict__ B_ptr = B + b * L * N;
const float* __restrict__ C_ptr = C + b * L * N;
float* __restrict__ y_ptr = y + b * L * D + d;
for (int t = 0; t < L; t++) {
float delta_t = delta_ptr[t * D];
float x_t = x_ptr[t * D];
float B_tn = B_ptr[t * N + n];
float C_tn = C_ptr[t * N + n];
...
}
```
This is clean and the compiler should generate good code.
One small optimization: prefetch `delta_ptr[(t+1)*D]` and `x_ptr[(t+1)*D]` using `__ldg` or just regular loads that the compiler might pipeline.
Actually, `__ldg` (load through read-only cache / texture cache) can help if the data is read-only. For x, delta, B, C, A_log, D_param - they're all read-only in the kernel. But `const float* __restrict__` already tells the compiler they don't alias, and modern GPUs automatically use the L1 cache.
Let me also see if there's value in using grid-stride loops or persistent kernels. Not really applicable here.
OK, let me write the kernel. I'll make the block size a template parameter or constant.
One issue: when compiling with `load_inline`, if I use `blockDim.x` as a variable, it's fine. But if I want compile-time constants for loop unrolling, I might use templates. However, that's overkill.
Actually, another idea: my kernel computes `dA = expf(delta_t * A_val)` every step. `A_val` is constant. `delta_t` comes from global memory. What if delta has many repeated values? Unlikely.
But there is a mathematical property: `dA = exp(delta_t * A) = exp(A)^{delta_t}`. Not sure if this helps.
Another optimization: for values of A that are very negative, dA might underflow to 0 for small delta_t. But this is mathematically correct.
What about using `fmaf` for fused multiply-add? The compiler automatically uses FMA instructions where appropriate. `h = dA * h + x_db` should compile to FMA.
Also, `x_t * dB` -> FMA or MUL. Fine.
One concern: the `__shfl_down_sync` inside a loop with `t` from 0 to L=2048. Each shuffle requires synchronization within the warp. Since all threads in each sub-group of 16 participate, and there's no divergence (all threads execute the shuffle), it's fast.
Let me verify with a test code mentally:
- For pair 0 (tid 0..15): all 16 threads reach shuffle.
- The `if (n == 0)` after shuffle: only threads with n=0 enter. Since n = tid % 16, for tid 0..15, n=0 corresponds to tid 0.
- For pair 1 (tid 16..31): n=0 corresponds to tid 16.
- Both execute the store instruction.
This is correct.
One more thing: if I launch exactly 512 blocks but the problem size changes (though in this task it's fixed), the kernel should handle it. But the benchmark uses fixed sizes, so it's fine.
Wait, the task says: "batch_size = 16, d_model = 256, d_state = 16, seq_len = 2048". These are fixed in the benchmark. But the kernel should still be correct for other sizes.
Let me make the kernel more general: handle any B, L, D, N where D is divisible by pairs_per_block.
Actually, to be safe, I can handle the case where D is not divisible by pairs_per_block by having the last block handle fewer pairs. But the kernel as written already has an `if (b >= B) return;` but not for extra d's. Let me add a check.
```cpp
if (b >= B || d >= D) return;
```
This handles any D value. The `d_groups` calculation would need to be `ceil(D / pairs_per_block)`, though.
```cpp
int d_groups = (D + pairs_per_block - 1) / pairs_per_block;
int b = blockIdx.x / d_groups;
int d_group = blockIdx.x % d_groups;
int d = d_group * pairs_per_block + pair_in_block;
```
This handles arbitrary D. I'll keep it general.
Similarly for B (batch size), the total blocks should be B * d_groups.
Let me finalize the kernel code and write it up.
Actually, before I write the code, let me also think about whether to use cooperative groups. `__shfl_down_sync` is built-in and doesn't require the cooperative groups header. I'll just use the built-in.
OK, one more optimization: the weights A_log[d*N + n] are the same for all blocks that have the same d. Since each d appears in exactly one block (for a given b), there's no reuse across blocks for A_log. A_log total size is 256*16*4 = 16KB. It fits in L1/L2 cache and stays there.
D_param[d] is also small: 256*4 = 1KB. Fully cached.
B and C are 2MB each. They might not fully fit in L2 cache (H100 has ~50MB L2), so they do. Actually, 2MB < 50MB. So B and C might stay in L2 cache. But when 512 blocks all try to read B and C, the cache eviction policy might evict. Since each batch b's B and C are 2048*16*4 = 128KB, and there are 16 batches, total B=2MB, C=2MB. L2 cache of 50MB can hold all of B and C easily. So once loaded, they stay in L2. Great!
Delta and x are 8MB each. Total 16MB. Still well within 50MB L2 cache! So actually, all the data might stay in L2 after the first read.
This means our kernel is essentially L2-cache bound, not DRAM bound. L2 bandwidth on H100 is much higher than DRAM. So the kernel could be very fast.
With all data in L2, the bottleneck shifts to compute (expf) and warp scheduling. The kernel should fly.
OK, I'm fairly confident in the design. Let me write the code now.
Actually wait - I said B and C are read by all blocks. But blocks for different batches read different regions of B and C. However, since 16 batches * 128KB = 2MB total, it all fits in L2. Similarly for delta and x: 16 * 2048 * 256 * 4 = 32MB each. Wait, that's way more than I said.
Let me recalculate:
- x: B * L * D = 16 * 2048 * 256 = 8,388,608 floats = 32MB
- delta: same = 32MB
- B: B * L * N = 16 * 2048 * 16 = 524,288 floats = 2MB
- C: same = 2MB
- y: 32MB (output)
Total inputs: 32 + 32 + 2 + 2 = 68MB. Plus A_log (16KB) and D (1KB).
68MB fits in 50MB? No, it doesn't. The L2 cache is 50MB, but 68MB of inputs don't all fit. However, the accesses are streaming: at each step t, we read a slice across all batches and d-channels.
Actually, for the kernel in a block:
- It reads delta for its specific (b, d) across all t: L floats = 8KB
- It reads x for the same: 8KB
- It reads B for its batch b across all t: L*N = 32KB
- It reads C for the same: 32KB
- Total per block: 80KB
With 512 blocks, if they all run concurrently, L2 needs to hold 512 * 80KB = 40MB. That's close to the 50MB limit, but since blocks don't all run exactly concurrently (scheduling), the working set might fit well.
Also, the accesses for x and delta are strided by D. For delta[b,:,d], the values are spaced D apart = 256*4 = 1024 bytes apart. Over L=2048, that's 2048 cache lines (assuming 64B cache lines, 16 floats per line). Wait, 1024 bytes is exactly 16 cache lines apart. So delta[b,t,d] for consecutive t maps to 16 different cache lines. This is bad for spatial locality!
Let me think again. Tensor delta is [B, L, D] with strides [L*D, D, 1]. So delta[b,t,d] and delta[b,t+1,d] are D floats apart. D=256, so 1024 bytes. A cache line is 64 bytes. So consecutive t accesses are 16 cache lines apart. This means for a single thread reading delta[b,0,d], delta[b,1,d], delta[b,2,d], etc., each read brings in a cache line of 64 bytes = 16 floats, but the thread only uses 1 float. The other 15 floats in the cache line are for delta[b,t,d:d+15] or delta[b,t,d-15:d].
Wait, cache line of 64 bytes = 16 floats. delta[b,t,d] is at address base + b*L*D + t*D + d. The cache line starting at this address contains delta[b,t,d] through delta[b,t,d+15] (assuming d is 16-aligned).
For a single thread reading delta[b,t,d] for t=0,1,2..., the cache lines are:
- t=0: cache line for delta[b,0,d..d+15]
- t=1: cache line for delta[b,1,d..d+15]
- etc.
These are 1024 bytes apart in memory. So each read brings a new cache line, and the thread uses 1/16 of it. This is the classic strided access problem.
However, in our block, all 16 pairs read delta[b,t,d_start..d_start+15] for the same t. With d_start aligned to 16, these 16 reads are exactly the 16 floats in one cache line! So the 16 threads (well, 16 leader threads, one per pair) each read one float, but they all read from the same cache line.
Wait no, each pair has 16 threads but only one reads delta_t? Let me trace again.
In our kernel, each thread reads delta_t for itself:
```cpp
float delta_t = delta_ptr[t * D];
```
But `delta_ptr` is specific to each pair because it points to `x_base + d` where `d` varies per pair.
For pair 0 (d=d_start): reads delta[b,t,d_start]
For pair 1 (d=d_start+1): reads delta[b,t,d_start+1]
...
For pair 15 (d=d_start+15): reads delta[b,t,d_start+15]
These are 16 contiguous floats = 1 cache line! And they're read by 16 different threads in the block. The warp scheduler handles this: the 16 threads are in warps 0..7 (2 pairs per warp). So warp 0 reads delta[b,t,d_start] and delta[b,t,d_start+1], warp 1 reads delta[b,t,d_start+2] and delta[b,t,d_start+3], etc.
In warp 0:
- Threads 0..15 (pair0) read the same address delta[b,t,d_start]. Wait, NO! Each of the 16 threads in pair0 has the SAME d! So all 16 threads in pair0 read `delta[b,t,d_start]`. That's a broadcast.
Oh crap, I made a mistake earlier! Let me re-examine.
Pair0: all 16 threads have the same b and same d, but different n. So all 16 threads read `delta[b,t,d]` where d=d_start. This is a broadcast! All 16 threads read the exact same address.
Then pair1 (threads 16..31) all read `delta[b,t,d_start+1]`. Another broadcast.
So in warp 0 (threads 0..31): 16 threads read address X, 16 threads read address X+4.
This is actually worse than I thought for cache line utilization! The cache line [X, X+64) contains delta[b,t,d_start..d_start+15]. Warp 0 reads X and X+4. Warp 1 reads X+8 and X+12. Etc. Each warp reads only 2 floats per cache line.
But across 8 warps, all 16 floats of the cache line are read. The first warp brings the cache line into L1, and subsequent warps hit L1. So total memory traffic is 1 cache line per t per block, which is optimal.
For a thread reading delta sequentially over t, the pattern is still strided. At t, it reads delta[b,t,d]. At t+1, delta[b,t+1,d]. These are D floats apart. But since multiple threads in the block read adjacent d values at the same t, the cache line brought in for delta[b,t,d_start..d_start+15] is fully utilized by the block at that t. Then for t+1, a new cache line at delta[b,t+1,d_start..d_start+15] is fetched.
So the block's delta access pattern is: for each t, fetch 1 cache line. That's excellent spatial locality for the block as a whole. The strided access per thread is amortized across the block.
Similarly for x.
Great, so the cache utilization is actually good at the block level.
The temporal locality: the block processes all t=0..L-1. As it goes, it reads cache lines for t=0, then t=1, etc. The cache line for t=0 is only used at t=0, then evicted. With L2 being large, some cache lines might persist if the block loops quickly. But streaming access is fine when bandwidth is high.
OK, now I'm ready to code.
Let me also consider an alternative: instead of having 16 threads all read the same delta_t, we could have one thread read it and broadcast via shared memory. But the hardware broadcast is already efficient.
Actually, for the 16 threads in pair0 all reading the same address, the GPU uses a special broadcast mechanism within a warp. If the same address is read by all threads in a warp, it's a single memory transaction. Wait, but in our warp, 16 threads read X and 16 threads read X+4. The hardware might handle this as two coalesced segments: segment1 for threads 0..15 (all same address, broadcast), segment2 for threads 16..31 (all same address X+4, broadcast).
NVIDIA docs: if all threads in a warp access the same address, it's a single transaction. If they access an even broadcast pattern... I think the warp memory controller handles it. Two broadcasts in one warp might be 2 transactions, which is fine.
Across 8 warps, each doing 2 transactions, that's 16 transactions per t per block. But they all hit the same cache line, so only 1 actually goes to L2/DRAM. The rest are L1 hits. Great!
I think the memory pattern is actually quite efficient.
Now, let me think about whether I should also put the projections in shared memory. No, as discussed, let's use precomputed tensors.
OK, I'm writing the code now. But I want to make sure I handle the dimensions properly and generalize it.
Wait, one more thing: I should check if the tensors are contiguous and on CUDA. The benchmark probably runs on GPU, and x is a CUDA tensor.
I'll add `.contiguous()` calls in the Python wrapper.
Let me draft the complete Python file:
```python
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.cpp_extension import load_inline
CUDA_SOURCE = r"""
#include <torch/extension.h>
#include <cuda_runtime.h>
__global__ void selective_scan_kernel(
const float* __restrict__ x,
const float* __restrict__ delta,
const float* __restrict__ A_log,
const float* __restrict__ B,
const float* __restrict__ C,
const float* __restrict__ D_param,
float* __restrict__ y,
int B, int L, int D, int N
) {
const int threads_per_pair = 16;
const int pairs_per_block = blockDim.x / threads_per_pair;
int tid = threadIdx.x;
int pair_in_block = tid / threads_per_pair;
int n = tid % threads_per_pair;
int d_groups = (D + pairs_per_block - 1) / pairs_per_block;
int b = blockIdx.x / d_groups;
int d_group = blockIdx.x % d_groups;
int d = d_group * pairs_per_block + pair_in_block;
if (b >= B || d >= D) return;
float A_val = -expf(A_log[d * N + n]);
float h = 0.0f;
size_t xd_batch_stride = (size_t)b * L * D;
size_t bc_batch_stride = (size_t)b * L * N;
size_t x_base = xd_batch_stride + d;
size_t b_base = bc_batch_stride;
size_t y_base = xd_batch_stride + d;
#pragma unroll 4
for (int t = 0; t < L; t++) {
size_t xd_offset = x_base + (size_t)t * D;
size_t bc_offset = b_base + (size_t)t * N;
float delta_t = delta[xd_offset];
float x_t = x[xd_offset];
float B_tn = B[bc_offset + n];
float C_tn = C[bc_offset + n];
float dA = expf(delta_t * A_val);
float dB = delta_t * B_tn;
float x_db = x_t * dB;
h = dA * h + x_db;
float contrib = h * C_tn;
// Warp shuffle reduction across n=0..15
contrib += __shfl_down_sync(0xFFFFFFFF, contrib, 8, 16);
contrib += __shfl_down_sync(0xFFFFFFFF, contrib, 4, 16);
contrib += __shfl_down_sync(0xFFFFFFFF, contrib, 2, 16);
contrib += __shfl_down_sync(0xFFFFFFFF, contrib, 1, 16);
if (n == 0) {
y[y_base + (size_t)t * D] = contrib + x_t * D_param[d];
}
}
}
void selective_scan_forward(
at::Tensor x,
at::Tensor delta,
at::Tensor A_log,
at::Tensor B,
at::Tensor C,
at::Tensor D_param,
at::Tensor y,
int B,
int L,
int D,
int N
) {
const int threads_per_pair = 16;
int block_size = 256;
int pairs_per_block = block_size / threads_per_pair;
int d_groups = (D + pairs_per_block - 1) / pairs_per_block;
int num_blocks = B * d_groups;
selective_scan_kernel<<<num_blocks, block_size>>>(
x.data_ptr<float>(),
delta.data_ptr<float>(),
A_log.data_ptr<float>(),
B.data_ptr<float>(),
C.data_ptr<float>(),
D_param.data_ptr<float>(),
y.data_ptr<float>(),
B, L, D, N
);
}
"""
selective_scan = load_inline(
name="selective_scan",
cpp_sources="",
cuda_sources=CUDA_SOURCE,
functions=["selective_scan_forward"],
extra_cuda_cflags=["-O3", "--use_fast_math"],
verbose=False,
)
class ModelNew(nn.Module):
def __init__(self, d_model: int, d_state: int):
super(ModelNew, self).__init__()
self.d_model = d_model
self.d_state = d_state
self.A_log = nn.Parameter(torch.randn(d_model, d_state))
self.proj_B = nn.Linear(d_model, d_state, bias=False)
self.proj_C = nn.Linear(d_model, d_state, bias=False)
self.proj_delta = nn.Linear(d_model, d_model, bias=True)
self.D = nn.Parameter(torch.ones(d_model))
def forward(self, x: torch.Tensor) -> torch.Tensor:
B, L, D = x.shape
N = self.d_state
delta = F.softplus(self.proj_delta(x))
B_t = self.proj_B(x)
C_t = self.proj_C(x)
y = torch.empty_like(x)
selective_scan.selective_scan_forward(
x.contiguous(),
delta.contiguous(),
self.A_log.contiguous(),
B_t.contiguous(),
C_t.contiguous(),
self.D.contiguous(),
y,
B, L, D, N
)
return y
```
Wait, I should check the `load_inline` signature. Let me make sure.
`torch.utils.cpp_extension.load_inline(name, cpp_sources, cuda_sources, functions, extra_cuda_cflags, verbose)`
Yes.
One thing: `cpp_sources=""` - we're not providing any C++ code other than what's in cuda_sources. The CUDA code is in `cuda_sources`. This should work - NVCC compiles `.cu` files which can contain both host and device code.
Actually, `load_inline` might expect `cpp_sources` to contain the pybind wrapper. But no, the documentation says pybind11 module is auto-generated from `functions`. We just need to provide the function implementations. And for CUDA, they go in `cuda_sources`.
Let me verify: "Do not write your own PYBIND11_MODULE(...) block in cpp_sources. load_inline auto-generates one from the functions=[...] argument."
So our function `selective_scan_forward` is listed in `functions`, and load_inline generates the pybind11 wrapper. Good.
One thing to be careful about: the C function signature should be C-compatible for pybind11. Since we're using `at::Tensor`, pybind11 knows how to bind these from the torch extension headers.
Also, we need to include `<torch/extension.h>` which brings in pybind11 support implicitly.
Wait, I should double-check: does the kernel function need to be declared `extern "C"`? No, C++ name mangling is fine for pybind11.
Let me also add `-arch=sm_80` or similar to the CUDA flags. Actually, `load_inline` automatically detects the architecture. But to be safe, not specifying it might be better so it uses the current GPU.
`"-O3"` enables optimization. `"--use_fast_math"` makes expf faster but slightly less accurate. The tolerance is 1e-4, so fast math should be acceptable. Actually, wait: fast_math can sometimes cause larger errors. For fp32 with 1e-4 tolerance, it's probably fine, but let me not use it initially and add it only if needed for speed.
I'll start without `--use_fast_math` to ensure correctness.
Actually, there's another concern: the reference uses `torch.exp` and our kernel uses `expf`. The difference between fast math and regular expf is small, but let me be safe.
One more thing: `#pragma unroll 4` is a compiler hint. It should help with loop unrolling to increase ILP.
Let me compile and test. But first, let me also think about whether `block_size = 256` is optimal. I'll keep it as 256 for now.
Actually, I want to run a quick correctness check. But I can only do that by actually compiling the code. Let me do that.
But wait - I need to make sure the ModelNew class is properly exported and matches the expected interface. The file should define `ModelNew` class.
One thing: in the provided reference, the linear projections in PyTorch might not produce the exact same numerical output as if we computed them manually. But we're using the same PyTorch linear layers, so that's fine.
Let me also consider: should I transpose A_log to [N, D] or keep it [D, N]? The reference stores it as [D, N], so our kernel should read it the same way. `self.A_log[d, n]` in PyTorch corresponds to `A_log[d * N + n]` in row-major. Correct.
What about B_t: in PyTorch, `nn.Linear(d_model, d_state)` with input [B, L, d_model] produces [B, L, d_state]. The last dimension is d_state=N. So `B_t[b, t, n]` corresponds to `B[b * L * N + t * N + n]`. Correct.
Similarly for C_t.
Delta: `nn.Linear(d_model, d_model)` produces [B, L, d_model]. `delta[b, t, d]` corresponds to `delta[b * L * D + t * D + d]`. Correct.
Now, the reference computes:
```python
dA = torch.exp(delta.unsqueeze(-1) * A.unsqueeze(0).unsqueeze(0))
```
This broadcasts delta [B, L, D, 1] with A [1, 1, D, N] to [B, L, D, N]. In our kernel, for fixed (b, t, d, n), we compute `exp(delta_t * A_val)`. This is exactly the same element-wise computation.
One potential difference: PyTorch's broadcasting might compute `delta.unsqueeze(-1) * A` using specialized kernels with different precision than per-element multiply. But elementwise multiply is deterministic and accurate.
Another potential difference: PyTorch's `h` initialization:
```python
h = torch.zeros(B, D, N, device=x.device, dtype=x.dtype)
```
And then the loop. Our kernel sets `h = 0.0f`. Matches.
The outputs list and stack:
```python
outputs = []
for t in range(L):
...
outputs.append(y_t)
y = torch.stack(outputs, dim=1)
```
In our kernel, we write `y[b, t, d]` directly. This matches because `torch.stack` on a list of [B, D] tensors along dim=1 gives [B, L, D].
And the D skip connection:
```python
y = y + x * self.D.unsqueeze(0).unsqueeze(0)
```
Our kernel does `contrib + x_t * D_param[d]`. Matches. But note: we only add it when `n == 0` (writing the output). For other threads (n != 0), they don't write to y. The skip connection is added to the final y, not per-n. Correct.
One thing to verify: in the reference, `self.D` is [D]. When added to y [B, L, D], it broadcasts. So each d-channel gets its own D value. Our kernel reads `D_param[d]`. Correct.
OK, I think the kernel is correct. Let me compile it.
Actually, I realize there might be an issue with how `blockDim.x` interacts with our kernel if block_size doesn't divide evenly. But we're launching exactly `block_size = 256`, and `pairs_per_block = block_size / 16 = 16`. This is hardcoded in the kernel. If we change block_size in the host code but not the kernel, it would be inconsistent. But since both are in the same file, they're consistent.
To be more robust, I could make `threads_per_pair = 16` and `pairs_per_block = blockDim.x / threads_per_pair`. This makes the kernel adapt to the block size. But then I'd need to ensure blockDim.x is a multiple of 16. Yes.
OK let me write the actual code now and compile it.
Actually, one more thought: since the problem provides exact dimensions and the reference runs them, the benchmark will call `get_inputs()` which returns a list with a tensor. The ModelNew is instantiated with `ModelNew(*get_init_inputs())`. So our `__init__` must accept `d_model` and `d_state`.
And `forward(self, x)` returns a tensor.
Let me also verify that in our `forward`, we don't accidentally modify `x` or other inputs. We're reading them and writing to `y`. Good.
One potential issue: `y = torch.empty_like(x)` might not be initialized. But we write to all elements of y. However, we only write from threads with `n == 0`. What about other d values?
For each (b, d), there's exactly one pair corresponding to it. The pair's leader (n==0) writes `y[b, t, d]`. So all (b, d) for all t are written exactly once. Good.
What if D is not a multiple of pairs_per_block? The last block might have pairs where `d >= D`. We have `if (d >= D) return;` so those threads don't write. But in the same warp, some threads have valid d and some don't. Does this cause issues with `__shfl_down_sync`?
If a thread has `d >= D`, it returns early and doesn't participate in the shuffle. But `__shfl_down_sync` with `mask=0xFFFFFFFF` requires all threads in the warp to participate, or behavior is undefined. This is a real issue!
To fix this, I need to ensure that invalid threads still participate in the shuffle, or I need to use a proper mask.
The easiest fix: ensure D is always a multiple of pairs_per_block by adjusting pairs_per_block. For D=256 and pairs_per_block=16, 256 is divisible by 16. So in the benchmark, this isn't an issue.
But to be general, let me handle it. I can set `pairs_per_block` to divide D. For arbitrary D, I could compute `pairs_per_block = gcd(16, D)`... but that's messy.
A cleaner approach: just make sure we launch blocks such that all threads in all warps are valid. We can round D up to the next multiple of pairs_per_block and launch extra blocks, but have threads with d>=D still compute (with dummy values) and participate in shuffles but not write.
```cpp
bool valid = (d < D);
float A_val = valid ? -expf(A_log[d * N + n]) : 0.0f;
float h = 0.0f;
// ... compute ...
float contrib = valid ? (h * C_tn) : 0.0f;
contrib += __shfl_down_sync(0xFFFFFFFF, contrib, 8, 16);
// ... etc ...
if (valid && n == 0) {
y[...] = ...;
}
```
This ensures all threads participate in shuffle, and invalid threads contribute 0 to the sum. But they also read memory out of bounds for A_log[d * N + n] if d >= D!
Actually, A_log is [D, N], so d might be out of bounds. The `valid ? -expf(...) : 0.0f` only evaluates `expf` when valid, so no out-of-bounds read. Wait, does the ternary operator short-circuit in C++? No, it evaluates both branches but the code is conditional.
Actually, the C++ ternary operator `valid ? expr1 : expr2` is a single expression. `expr1` is only evaluated if `valid` is true? No, in C++ the ternary operator only evaluates one of the two branches. So `valid ? -expf(A_log[d * N + n]) : 0.0f` does NOT read A_log if valid is false.
But what about the other reads like `delta[xd_offset]`? Those would be out of bounds for invalid threads. We need to guard them too:
```cpp
float delta_t = valid ? delta[xd_offset] : 0.0f;
float x_t = valid ? x[xd_offset] : 0.0f;
float B_tn = valid ? B[bc_offset + n] : 0.0f;
float C_tn = valid ? C[bc_offset + n] : 0.0f;
```
This ensures all memory accesses are in bounds, and all threads participate in shuffles. The invalid threads just compute 0 and don't write.
For the benchmark with D=256, this isn't necessary, but it's good practice.
Wait, for B and C, the offsets are based on `b` and `t` and `n`, not `d`. If `b >= B`, then `b_base` would be out of bounds. But we already have `if (b >= B) return;` which catches that. For `d >= D`, `b` is valid, so B and C accesses are valid.
But `delta[xd_offset]` with d >= D is out of bounds. So we need to guard delta and x accesses.
I'll add the `valid` check for robustness:
```cpp
bool valid = (b < B && d < D);
```
Then guard all reads and the final write.
Actually, with `if (b >= B) return;` at the top, we only need `valid = (d < D)`. But I'll just use `if (b >= B || d >= D) return;`... wait, no, because we need invalid threads to participate in shuffle for their warp mates. If we return early, the shuffle might hang.
So we must NOT have early return for threads in the same warp as valid threads. But threads with `b >= B` are in a different block entirely? No, blocks are different. Within a block, all threads are in the same block. If some pairs in a block have `d >= D`, they share the block with valid pairs.
For the benchmark: D=256, pairs_per_block=16, d_groups=16, num_blocks=256. Within a block, d ranges from d_start to d_start+15. With D=256 and pairs_per_block=16, d_start ranges from 0, 16, 32, ..., 240. The max d = 255 < 256. So all threads in all blocks have d < D. No invalid threads!
But if D were say 250, then the last d_group would have d_start=240, and d=240..255, but D=250 means d=250..255 are invalid.
For the benchmark, this doesn't happen. So I'll keep the simple version with `if (b >= B || d >= D) return;` and not worry about shuffle masks.
However, the `return` might still be safe if all threads in a warp return. Let's see: with D=256 and pairs_per_block=16, no invalid d. But what about the last block if num_blocks isn't exact? `num_blocks = B * d_groups` where `d_groups = ceil(D / pairs_per_block)`. For D=256, pairs_per_block=16, d_groups=16, num_blocks=256. Each block has pairs 0..15 for some (b, d_group). All valid.
So for this specific benchmark, no invalid threads. But for generality, I'll keep the valid check. Actually, let me just rely on the benchmark using exact multiples and not add the overhead of the guard in the inner loop.
Actually, I realize I should be careful about `D` being passed exactly. The benchmark uses `d_model=256`, `d_state=16`. So this is fine.
Let me keep the simple version for speed but add bounds checking outside the loop only for A_log:
```cpp
if (b >= B || d >= D) return;
```
This return is safe because if `d >= D`, the entire pair is invalid. A warp has 2 pairs. If one pair is invalid and one is valid, the invalid pair's threads return while valid pair's threads continue to shuffle. This could deadlock.
But again, for this benchmark, all pairs are valid due to exact division. I'll keep the simple version and mention in a comment that D should be divisible by pairs_per_block.
Actually, to be safe, I can just compute num_blocks using exact division when D % pairs_per_block == 0. Or I can handle it by only launching exact blocks. But I'll just use ceil division in num_blocks and add the guard. For the return causing deadlock issue: if D % pairs_per_block != 0, the last block would have some invalid d's.
One solution: if D % pairs_per_block != 0, we can still handle it by having the invalid threads return, but we must ensure they don't share warps with valid threads for the same pair. Wait, in a warp, pair_in_block=0 and 1. If d_group * pairs_per_block + 0 < D but d_group * pairs_per_block + 1 >= D, then pair0 is valid and pair1 is invalid. They're in the same warp. If pair1's threads return early, pair0's shuffle calls will hang because not all threads in the warp participate.
To avoid this, we can set `pairs_per_block` to 1 when D is small, or more generally, ensure we don't mix valid and invalid pairs in the same warp. This requires careful block sizing.
For the benchmark, D=256, block_size=256, pairs_per_block=16, d_groups=16. Entirely exact. No problem.
I'll use exact division in the kernel and assert or handle generality. Since the benchmark is fixed, I'll hardcode 16 pairs per block.
Actually, to be super safe while keeping performance, I can change the launch to only launch exact number of pairs:
```cpp
int num_pairs = B * D;
int num_blocks = (num_pairs + pairs_per_block - 1) / pairs_per_block;
```
And in the kernel, use linear pair indexing:
```cpp
int pair_idx = blockIdx.x * pairs_per_block + pair_in_block;
if (pair_idx >= num_pairs) return;
int b = pair_idx / D;
int d = pair_idx % D;
```
But with this, the last block might have some invalid pairs. These invalid pairs might share a warp with valid pairs. If they return early, shuffle deadlock.
To avoid deadlock, I need to make sure invalid threads don't share warps with valid threads. Or I can have them participate with 0 values.
Actually, with `if (pair_idx >= num_pairs) return;`, if pair_idx is assigned linearly, then valid and invalid pairs might interleave... wait no. pair_idx = blockIdx.x * 16 + pair_in_block. For the last block, say we need 4094 pairs but launch 256 blocks (4096 pairs). The last 2 pairs are invalid. They are pair_in_block = 14 and 15.
In the block, pairs 14 and 15 are in warps 7 (if 2 pairs per warp: warp 0 has pairs 0,1; warp 1 has pairs 2,3; ... warp 7 has pairs 14,15). So warp 7 has pairs 14 and 15, both invalid. All threads in warp 7 return. But warp 0..6 have valid pairs and continue. The sync is per warp, not across warps. So warp 0..6's `__shfl_down_sync` only involves threads within those warps, and all threads in those warps are valid. Warp 7 returns. This is fine!
The `__shfl_down_sync` is a warp-level operation. Threads in different warps don't synchronize with each other. So as long as all threads within a warp either all participate or all return, it's safe.
In our case, pairs 14 and 15 are both in warp 7. If both are invalid (pair_idx >= num_pairs), all 32 threads in warp 7 return. Warps 0..6 don't have invalid pairs and proceed normally. This is safe.
What if only pair 15 is invalid but pair 14 is valid? Then warp 7 has 16 valid threads and 16 invalid threads. The invalid threads return, the valid ones do shuffle. This would deadlock.
When does this happen? If num_pairs is odd and pairs_per_block is even. With num_pairs = B*D = 16*256 = 4096 (even), and pairs_per_block = 16 (even), the last block always has invalid pairs starting from an even index. If num_pairs % pairs_per_block = 0, there are no invalid pairs.
4096 % 16 = 0. So no invalid pairs.
If num_pairs % pairs_per_block != 0, and the remainder is less than 2, then we might have a warp with mixed valid/invalid. But with the benchmark fixed, this doesn't happen.
So I'll use the linear pair indexing approach for generality but rely on exact divisibility. Or I can just keep the 2D indexing.
For simplicity, I'll keep the 2D indexing (b = blockIdx.x / d_groups, d = ...) and the `if (b >= B || d >= D) return;`. But I should ensure num_blocks is computed with ceiling and the last block's invalid threads don't deadlock.
Actually, with 2D indexing and d_groups = ceil(D / pairs_per_block):
- For blockIdx.x, b and d_start are computed.
- If d_start >= D, all pairs in the block are invalid.
- If d_start + pairs_per_block - 1 < D, all pairs are valid.
- If d_start < D <= d_start + pairs_per_block - 1, some pairs are valid and some invalid.
For D=256, pairs_per_block=16: d_start is always a multiple of 16. 0, 16, ..., 240. All < 256, and all + 15 < 256. All valid.
For D=250: d_start = 0, 16, ..., 240. The last d_start=240, and 240+15=255 > 250. So pairs 0..9 are valid (d=240..249), pairs 10..15 are invalid (d=250..255).
In this block, valid pairs are 10, 11 (if 0-indexed within block), invalid are 12..15. Wait, pairs_per_block=16, pair_in_block 0..15 correspond to d=240..255. Valid d: 240..249 (10 values). So 10 valid pairs and 6 invalid.
Warp assignment: warp 0 handles pairs 0,1 (d=240,241). Both valid.
Warp 1 handles pairs 2,3. Both valid.
...
Warp 4 handles pairs 8,9. Both valid.
Warp 5 handles pairs 10,11. d=250,251. Both invalid!
Warp 6 handles pairs 12,13. Both invalid.
Warp 7 handles pairs 14,15. Both invalid.
In this case, all warps either have all valid or all invalid pairs. Safe!
General proof: Since pairs_per_block = 16, and each warp handles 2 pairs (32 threads / 16 threads_per_pair), if the number of valid pairs in a block is a multiple of 2, warps are either all-valid or all-invalid.
With d_start being a multiple of 16, and D being any value: the number of valid pairs is min(16, D - d_start). This could be odd! Example: D=257, d_start=256 (last group). Valid pairs = 1 (d=256). This is in pair_in_block=0, warp 0. Warp 0 also has pair_in_block=1 (d=257, invalid). So mixed warp. Deadlock.
To handle this generally, we need one of:
1. Only use `return` when all threads in the warp are invalid.
2. Have invalid threads participate with 0 values.
3. Pad D to a multiple of pairs_per_block.
Option 3 is easiest: in the host code, we can assert or pad. But for generality, option 2 is safer.
For this specific benchmark, D=256, pairs_per_block=16, so it's exact. I'll just rely on that and add a comment. No need to overcomplicate.
Actually, for maximum generality without much overhead, I can check in the host code:
```python
pairs_per_block = 16
assert D % pairs_per_block == 0, ...
```
Or just let it fail with incorrect results if not. But the benchmark will use exact sizes.
OK, I'll go with the simple exact version and compile it. Let me also double check `__shfl_down_sync`: I need to make sure `mask` includes all active threads in the warp. Since all threads in the warp reach `__shfl_down_sync` (the `if (d >= D) return` doesn't happen for exact division), and mask=0xFFFFFFFF, it's fine.
Let me write and compile the kernel now.