We need to write a custom CUDA kernel to compute row-wise softmax over CSR edge scores, for multiple attention heads. The reference loops over nodes, each node corresponds to a contiguous slice of edge_scores defined by row_ptr, and computes softmax along dim=0 (over edges for that node) independently per head. We need to define ModelNew with load_inline. We must match constructor (no args). forward accepts row_ptr (int32 cpu or maybe cuda? Could be on CPU; but likely same device as scores; to be safe we can operate on device of edge_scores). row_ptr might be on CPU because used for indexing; but in reference they treat row_ptr as regular tensor, use .item() repeatedly (which is CPU). However in typical usage, row_ptr may be on CPU. The custom kernel could accept row_ptr as a cpu tensor or move to device as needed. But for performance, copying row_ptr to device once is okay. But note that if row_ptr stays on CPU, doing per-node launch from CPU is okay because there are only 384 nodes. Simpler: we can just copy row_ptr to the same device as edge_scores, and then have a CUDA kernel per edge or per node using row_ptr. Goal: speedup over reference, which is a slow Python loop with per-row softmax. 384 rows, avg degree 18, num_heads 8. So about 384*18 ~ 6912 edges, 8 heads. That's tiny. But kernel overhead can dominate if not careful; however the reference does Python loop with .item() and torch.softmax calls, each launching a kernel. So likely a single custom kernel that processes all rows will be much faster. Even with such small data, reduction per row is small but overhead reduction helps. Since there are 384 rows, we could launch one thread per row and each thread loops over its edges and per-head compute max, exp, sum, then normalize. That avoids parallelization inside the row (which might be okay because avg degree ~18). But we need to account for varying row lengths (clamped 1 to ~20). A block size like 256 threads, each processing one node. Each thread loops over edges for that node. Because num_heads=8, each thread can handle the 8 heads in a loop (or unrolled). Might be efficient enough. Alternatively we could assign each warp to a node and parallelize over edges within warp, then do warp reduce. But with 18 edges, warp reduce may not be much benefit and added complexity. Since node count is small, one thread per node is fine. Let's design kernel: Input: - const int* row_ptr (size num_nodes+1) - const float* edge_scores (num_edges * num_heads) - float* out (same size) - int num_nodes - int num_heads CUDA kernel: __global__ void csr_softmax_kernel(const int* row_ptr, const float* edge_scores, float* out, int num_nodes, int num_heads) { int node = blockIdx.x * blockDim.x + threadIdx.x; if (node >= num_nodes) return; int start = row_ptr[node]; int end = row_ptr[node + 1]; // For each head for (int h = 0; h < num_heads; ++h) { float max_val = -1e30f; // compute max for (int e = start; e < end; ++e) { float val = edge_scores[e * num_heads + h]; if (val > max_val) max_val = val; } float sum = 0.0f; for (int e = start; e < end; ++e) { float val = edge_scores[e * num_heads + h]; sum += expf(val - max_val); } float inv_sum = 1.0f / sum; for (int e = start; e < end; ++e) { float val = edge_scores[e * num_heads + h]; out[e * num_heads + h] = expf(val - max_val) * inv_sum; } } } That seems straightforward. But we must match reference numerically within fp32 tolerance 1e-4. Reference uses torch.softmax which is numerically stable: subtracts max, exp, divide by sum. Our kernel does same with float, should match. Potential issue: edge_scores may be large causing overflow? With max subtraction, safe. Potential issue: row_ptr might be int32 on CPU. In load_inline wrapper, we can ensure row_ptr is on same device. If row_ptr is on CPU, we need to copy to CUDA and maybe also convert to int32 (it already is). Let's handle in forward: move row_ptr to edge_scores.device. But also row_ptr may already be CUDA? The get_inputs returns on CPU (default). So we need to copy. One more thing: `torch.utils.cpp_extension.load_inline` at module load time. Need to write the code accordingly. We should also consider using `__shared__` but not needed. We need to compile, test. Let's code it. Wait, we need exactly `class ModelNew` defined. Also `get_inputs` and `get_init_inputs` already defined externally. We just provide ModelNew. Let's write kernel. Important: load_inline parameters: `cuda_sources` contains the CUDA code and C++ wrapper. `cpp_sources` maybe empty. `functions` list wrapper names. We need to write a host function that takes Tensor objects, calls the kernel. For simplicity: void csr_softmax(torch::Tensor row_ptr, torch::Tensor edge_scores, torch::Tensor out) { int num_nodes = row_ptr.numel() - 1; int num_heads = edge_scores.size(1); const int* row_ptr_d = row_ptr.data_ptr<int>(); const float* scores_d = edge_scores.data_ptr<float>(); float* out_d = out.data_ptr<float>(); int threads = 256; int blocks = (num_nodes + threads - 1) / threads; csr_softmax_kernel<<<blocks, threads>>>(row_ptr_d, scores_d, out_d, num_nodes, num_heads); } Need to include <torch/extension.h>. The wrapper is in `cuda_sources` or `cpp_sources`? We can put both in `cuda_sources` or split. Typically we put the CUDA kernel in `cuda_sources`, and the host wrapper also in `cuda_sources`. `load_inline` with `cuda_sources` will compile with nvcc. `cpp_sources` can be empty. `functions=['csr_softmax']`. Let's verify syntax. Since row_ptr is int32, ensure we use .data_ptr<int>() but return type is int32_t maybe. Use int32_t. Let's draft Python file. Potential performance: launch overhead? blocks = 2 (for 384 nodes, 256 threads). That's fine. Each thread loops over up to ~20 edges * 3 passes * 8 heads = ~480 operations. Very small but okay. Now compile and test. Let's produce the code and call compile_kernel.

Let me fix the `load_inline` call - I need to provide `cpp_sources` argument. I can provide an empty string for `cpp_sources` if I only have CUDA code.

Basic kernel compiled and passed correctness. But it's quite naive — one thread per node which will lead to poor occupancy and warps with very imbalanced workloads (since some rows have 16, some have 20 edges). We should design a more efficient kernel. We have avg_degree around 18, num_heads = 8, num_nodes = 384. The input is float32. Total edges ~ 18*384 = 6912. Each row is independent softmax per head, across edges. Since each row has ~18 edges and 8 heads, total elements per row: ~144. That's quite small. A warp per row could work, processing each head perhaps cooperatively. Given that the degree is small (~16-20), a block per row is too much. A warp per row is good. 32 threads per row, but only 18 elements to read. We could assign 1 warp per row. For each head, use warp-level reduction to compute max and sum. Alternatively, a block processes multiple rows e.g. 256 threads = 8 warps, each warp handles one row. That would process 8 rows per block. With 384 nodes, that's 48 blocks. On an H100, this is fine but not saturating the GPU. Still, speedup over CPU-like torch loop is expected. Given the small problem size (384 nodes), the loop on CPU already is fast but Python overhead + kernel launch + torch.softmax for each row is large. Even the naive kernel should give significant speedup. But we can do better. First, let me try the naive kernel with submit to see if speedup is enough. But maybe I should optimize first. Let me try a warp-per-row kernel. Use 32 threads per warp. For each row, each warp: - For each head h (0..7): - Load edge_scores for edges in strided manner from start to end. - Compute warp max. - Compute warp sum of exp. - Each thread writes its assigned edges. Since num_heads=8 and degree~18, each warp does 8 iterations, each with ~18 loads and stores. Total loads per warp = 8*18=144. Each thread does ~144/32 = 4.5 loads. Very light. But better than 1 thread per row doing all sequential loads. Also, shared memory may not be needed; we can keep values in registers while doing the load twice (once for max, once for sum). However, degree is small, so we could load into registers or shared memory. Since max degree is clamped from avg 18 plus/minus 2, max degree is 20. 20*8 = 160 floats = 640 bytes per row. Not much. But registers: each thread in a warp handling a row could store ~ceil(20/32)=1 value? Not enough. Actually if we want to reuse across max and sum passes, we need to either store in shared memory or reload. Reloading is fine since it's from L1 and only 18 elements * 8 heads = 144 floats = 576 bytes per row. That's tiny. So the warp-per-row kernel reloading is fine. But to reduce the number of passes, we could load once into shared memory for the row, then do max and sum from shared memory. Let's implement a block-per-multiple-rows kernel. Each block has WARPS_PER_BLOCK warps. Each warp handles one row. We can have e.g. 8 warps per block (256 threads). So each block handles 8 rows. Shared memory: For each warp (row), we allocate MAX_DEGREE floats per head? Actually we can allocate per warp a small shared memory buffer. Since max degree is around 20, and we have 8 warps, we need 8 * 20 * 8 heads? If we store per edge per head, but we can process heads sequentially on the fly. Let's just use registers and reload. Simpler. Implementation details: - grid: (num_nodes + WARPS_PER_BLOCK - 1) / WARPS_PER_BLOCK blocks. - block: WARPS_PER_BLOCK * 32 threads. - warp_id = threadIdx.x / 32 - lane = threadIdx.x % 32 - row = blockIdx.x * WARPS_PER_BLOCK + warp_id - if row >= num_nodes: return - start = row_ptr[row]; end = row_ptr[row+1]; - for h in 0..num_heads-1: - Compute max_val across edges [start, end) for this head. - Use __shfl_sync for warp reduce, or use WarpReduce from cub. For simplicity, use __shfl_sync. - Then compute sum of expf(val - max) across edges. - Then write back expf(val - max) / sum for each edge. Warp reduction for max and sum. We need each thread to iterate over edges in strided fashion by lane id, and then use warp-level reduction. Pseudo: float max_val = -1e30f; for (int e = start + lane; e < end; e += 32) { float v = edge_scores[e * num_heads + h]; max_val = fmaxf(max_val, v); } // warp reduce max #pragma unroll for (int offset = 16; offset > 0; offset /= 2) max_val = fmaxf(max_val, __shfl_down_sync(0xFFFFFFFF, max_val, offset)); // broadcast max max_val = __shfl_sync(0xFFFFFFFF, max_val, 0); float sum = 0.0f; for (int e = start + lane; e < end; e += 32) { sum += expf(edge_scores[e * num_heads + h] - max_val); } // warp reduce sum #pragma unroll for (int offset = 16; offset > 0; offset /= 2) sum += __shfl_down_sync(0xFFFFFFFF, sum, offset); sum = __shfl_sync(0xFFFFFFFF, sum, 0); float inv_sum = 1.0f / sum; for (int e = start + lane; e < end; e += 32) { out[e * num_heads + h] = expf(edge_scores[e * num_heads + h] - max_val) * inv_sum; } This is good. But we re-read edge_scores three times per head. 3 * 8 * ~18 = 432 loads per warp. Still very small. But we could reduce to 2 reads if we compute expf and store in a register array? Not feasible due to variable degree. But maybe use shared memory: load all edges for this head into a shared buffer, then do max reduction on that buffer, then sum, then write back. However, that requires synchronization within warp, but warps are SIMT so they execute together; we can use shared memory per warp. Since degree is small, we could have each thread load one or zero edges into shared mem, then do reductions. But the strided loop with shfl is simpler and likely fast enough for this size. But is this faster than the naive kernel? The naive kernel has one thread per node doing sequential access. With 384 threads total, many warps will be active but each warp has only 1 active thread (others idle). That's bad for occupancy but also memory access patterns are not coalesced because each thread reads contiguous floats, but within a warp only thread 0 reads, so memory is fetched in 32-byte lines but mostly wasted. With warp-per-row, all 32 lanes participate and reads are coalesced (consecutive threads read consecutive edges). This gives 32x better memory bandwidth utilization. Also, warp-level reduction avoids shared memory and is fast. Let's test this kernel. One consideration: The row_ptr in the reference is created on CPU, so passed to forward may be CPU tensor. We move it to GPU inside forward. Need to ensure the kernel receives all tensors on GPU and contiguous. Also, we might want to process multiple rows per warp? Since degree is ~18, one warp per row is slightly underutilized (only 18/32 lanes have work per iteration). But since there are 8 heads, total work per warp is ~144 elements. That's okay. Could process 2 rows per warp? Then each lane handles edges from two rows, which complicates indexing and reduction. Not worth it. Alternatively, process 1 thread per head per row? Since num_heads=8, a warp of 32 threads could handle 4 rows simultaneously (8 threads per row). But reduction across 8 threads is a bit more complex but can be done with shfl. However, memory coalescing: if thread 0..7 handle row 0 edges 0.., they would access edge_scores[e*8 + h] where h is fixed per thread. So within a row, consecutive threads access consecutive head dimension. That is actually coalesced if we arrange properly: thread h accesses head h. But e varies across iterations. Not sure it helps much. I think warp-per-row is the best simple optimization. Also, for even more performance, we could use vectorized loads (float4) if degree >= 4. But degree is ~18, so float4 could help, but the indexing is [e * num_heads + h]. Since num_heads=8, consecutive e mean offset 8 floats. That is 32 bytes, which is one float4. So we can load float4 for 4 consecutive edges? Actually edge_scores[e*8 + h], if we fix h, consecutive e are offset by 8 floats = 32 bytes. So a float4 load would load 4 edges at once (e, e+1, e+2, e+3) for the same head. That is perfectly aligned if start is aligned and e increments by 4. Since the row lengths are arbitrary, alignment may not be guaranteed. But on GPU, misaligned float4 loads can still be okay but may split into multiple transactions. To keep it simple, I'll skip vectorization. Another optimization: Process multiple heads in a single loop to reuse loaded values? Not possible since each head is independent. Also, degree varies per row (16 to 20). The loop `for (int e = start + lane; e < end; e += 32)` will have at most 1 iteration because degree < 32. So each lane either loads 1 element or none. That's fine. Let's write the kernel. gridDim.x = (num_nodes + WARPS_PER_BLOCK - 1) / WARPS_PER_BLOCK blockDim.x = WARPS_PER_BLOCK * 32 WARPS_PER_BLOCK can be 8 for 256 threads. Let's compile and test.

The warparized kernel passed correctness. Numerical precision is still excellent (max_abs_err 1.79e-07, max_rel_err 3.91e-07, well within 1e-4). I should now consider whether warp-per-row is optimal or if thread-per-row (sequential) might be faster given the small degree (average ~18). Let me compute occupancy and think about memory access. Degree is small (~18), num_heads = 8, so each warp does 18*8 = 144 iterations of coalesced reads. With 8 warps per block (256 threads), blocks process a chunk of 8 rows. We need blocks = ceil(384/8) = 48 blocks. This is plenty, and should occupy many SMs. But sequential per thread might actually have better data locality since the degree is so small. However, with small degree, warp-parallel reduction does very little useful work (only reducing across ~18 values but with 32 threads). The instruction overhead of warp shuffle might dominate. Actually, one more thought: I can further optimize by unrolling the loop over heads and using vectorized loads if possible. But the memory layout (num_edges x num_heads) with num_heads=8 means each row's head data is contiguous in memory. So within a row, for a given edge e, the 8 floats from edge_scores[e*8 : e*8+7] are contiguous. I can use float4 to load 4 at a time, or float2. But that complicates the max+sum+softmax logic. For simplicity and given the relatively small problem, the warp kernel is likely sufficient to beat PyTorch. Actually, wait — I should consider whether the PyTorch reference is CPU bound or GPU bound. The reference uses `torch.softmax(..., dim=0)` which is a highly optimized PyTorch kernel. Since each row is small (18 elements x 8 heads), launching 384 separate softmax calls in a Python for loop is the slow part. Our single fused CUDA kernel eliminates Python overhead and kernel launch overhead, so even a simple kernel should be much faster. Given that correctness passes and the design is reasonable, I think I should submit the warp kernel. But let me also check if there are any hidden issues by running static_check. Actually, let me also consider if I can improve performance even more by reducing redundant expf calls. In the current kernel, I compute exp twice per element (once for sum, once for output). I can precompute and store in shared memory, but that requires more shared memory. With degree ~18, 18*8 floats = 576 bytes, very small. Let me make a version that stores intermediate exp values in shared memory to avoid recomputing. But honestly, this is a level 1 problem and the speedup over the reference (Python loop) should be huge regardless. Let me just static_check and submit.