Let me analyze the problem: 1. We have a multi-head row-wise softmax over CSR edge scores 2. The input `row_ptr` defines the CSR structure (row_ptr[i] to row_ptr[i+1] gives the edge indices for node i) 3. `edge_scores` has shape (num_edges, num_heads) 4. For each node, we need to apply softmax across its edges for each head independently The reference implementation loops over nodes in Python, which is very slow. I need to write a custom CUDA kernel that parallelizes this. Key observations: - num_nodes = 384 - avg_degree = 18 - num_heads = 8 - Total edges ≈ 384 * 18 = 6912 The softmax is applied per-node per-head. So for each node i and head h: - Find edges from row_ptr[i] to row_ptr[i+1] - Compute softmax of edge_scores[start:end, h] This is essentially a segmented softmax problem. For CUDA optimization: 1. We could assign one thread per (node, head) pair, but that might not be enough parallelism 2. Better: one warp or block per node, parallelizing across edges and heads 3. Even better: use parallel reduction for softmax computation Given the small sizes (384 nodes, ~18 edges per node, 8 heads), we could: - Use one thread block per node, with multiple threads handling edges and heads - Use warp-level primitives for efficient reduction Let me design the kernel: Option 1: One thread per (node, edge, head) - but we need to compute max and sum across edges, which requires reduction Option 2: One block per node, use shared memory for reduction across edges Option 3: One warp per (node, head) pair Given avg_degree = 18, a warp (32 threads) can handle one head's edges per node. We'd have 384 * 8 = 3072 warps, which is good parallelism. For each warp handling (node, head), I need to find the max value across edges, compute exp(score - max), sum the exponentials, then normalize. This follows the standard numerically stable softmax pattern with max subtraction before exponentiation to prevent overflow. With 18 edges per node on average, a single warp can efficiently manage the edges for one head. The warp processes edges in strides, with each thread computing the max, sum, and normalization for its assigned edges. If edges exceed the warp size, threads loop through multiple iterations. After processing all edges, a warp reduction combines thread results, then each thread normalizes its edges by the final sum. I'm implementing a softmax kernel where each warp handles one (node, head) pair, using bitwise operations to extract lane IDs and warp count, then iterating through edges in chunks of 32 to find the maximum value across all edges for numerical stability. Within the loop, I check bounds against the edge count, load scores into registers, and track the local maximum. After processing all edges, I perform a warp-level reduction using XOR operations to combine all thread maxima into a single value that's broadcast across the warp. Next, I iterate through the edges again, this time computing exp(score - max) for each element and accumulating the sum. I apply the same warp reduction pattern to combine these sums across threads, giving me the total for normalization. For the final output, I loop through the edges a third time, normalizing each value by dividing by the total sum computed in the previous step and storing the result. The kernel needs optimization for the GPU architecture. With 384 nodes and 8 heads, I have 3072 warps total, which translates to roughly 96 blocks of 32 warps each on a typical GPU. The per-warp work involves loading and storing edges twice while performing reductions, which should stay efficient for small average degrees. I could consider using one thread per (node, edge, head) tuple with atomic operations, but that introduces contention. Instead, I'll try assigning one thread block per node and use shared memory for coordinated reductions across all heads and edges in that node. In this block-per-node approach, I load edge scores into shared memory, compute the maximum score and exponential sum across all threads in the block, then normalize values. With 384 nodes and 8 heads yielding around 6912 total edges, distributing work across blocks with shared memory coordination should handle the workload effectively. _max, then compute exp(score - max) and sum these values across threads. The kernel uses block reduction to efficiently aggregate the sum, then each thread normalizes its score by dividing by the total sum. Finally, the normalized scores are written back to the output array. The second approach has advantages: shared memory enables efficient reduction across threads, a single pass handles all heads simultaneously, and the reduction operations themselves run quickly on GPU hardware. However, there's a fundamental problem. With an average degree of 18 and 8 heads per node, each block would need roughly 144 threads—but not all nodes have uniform degrees. The block size must accommodate the maximum degree across all nodes, yet we can't determine this at compile time. Using a fixed block size of 256 threads could work, but threads beyond the actual edge count would remain idle, wasting resources. The original single-thread-per-(node, edge, head) design avoids this issue. Each warp handles one (node, head) pair, which aligns well with the problem structure and doesn't require shared memory coordination across threads. For edge traversal, I'm using a stride pattern where each thread in a warp processes edges at intervals of 32 (the warp size). Thread 0 takes edges at indices 0, 32, 64..., thread 1 takes edges 1, 33, 65..., and so on. This ensures all edges for a given node get processed across the warp. In the first pass, I find the maximum score across all edges in the range. Each thread tracks its local max while iterating through the edges, then a warp reduction combines these local maxima into a single global max for the node. This max is needed for numerical stability in the softmax computation. _shfl_xor_sync(0xffffffff, local_sum, offset); \ } \ } while(0) But actually, I'm overcomplicating this. Let me simplify and just use the correct __shfl_xor_sync directly. For WarpReduceMax, each thread participates by XORing its value with progressively larger offsets (16, 8, 4, 2, 1) to combine results across the warp, keeping the maximum each time. WarpReduceSum works similarly for summing values across threads. Then in the kernel itself, each warp handles one (node, head) combination. I need to extract the node and head indices from the warp ID, where head = warp_id % num_heads and node = warp_id / num_heads. If the node index exceeds the total nodes, I return early. Otherwise, I locate the edge range for that node using row pointers, calculate the edge count, and iterate through the edges to find the maximum score value. Each thread processes a subset of edges based on its lane ID, then the warp performs a reduction to consolidate the maximum across all lanes. Next, I compute the sum of exponentials by iterating through the edges again. For each edge, I subtract the maximum for stability, apply the exponential function, and accumulate the sum locally before reducing across the warp. Then I normalize the scores by dividing each exponential value by the total sum and storing the result. Finally, I calculate the total number of warps needed as (num_nodes * num_heads) and launch the kernel with a block count derived from dividing total_warps by the warp size. Each CUDA block contains multiple warps, so for 3072 total warps with a block size of 256, I get 12 blocks, which is insufficient for full GPU utilization. Using 128 threads per block instead gives me 24 blocks, still not ideal. I should consider restructuring to have one block per node with threads distributed across edges and heads for better parallelism. Actually, the warp-per-(node,head) approach remains sound since warps execute efficiently across different SMs. With 3072 independent warps and 6912 total elements, each handles roughly 2.25 elements on average. For the implementation itself, I need to properly construct the string containing the CUDA code, then load it as a C++ extension with the necessary include directives and namespace declarations. The kernel uses warp-level primitives to compute softmax over CSR edge scores. Each warp processes one (node, head) pair, identifying the maximum value across edges using a reduction pattern, then normalizing the scores. The edge indices are determined from the row pointers, and threads within the warp cooperate to handle multiple edges per row. After finding the maximum, the kernel computes the sum of exponentials by iterating through edges again, applying exp(score - max) for numerical stability. A warp reduction accumulates these partial sums. Finally, each thread normalizes its assigned edges by dividing the exponential values by the total sum and writing the results to the output array. The launch configuration computes total warps as the product of nodes and heads, then divides by 32 to determine block count. With 3072 warps and 128 threads per block, this yields 24 blocks—relatively modest but functional. For improved GPU utilization, I could restructure to increase the block count. However, there's a critical issue: the average degree calculation is off. With 384 nodes and average degree of 18, that's 6912 total edges, which divided across 3072 warps gives only about 2.25 elements per warp on average. The problem setup with num_nodes = 384, avg_degree = 18, and num_heads = 8 confirms this mismatch in the parallelism breakdown. For each node-head pair, there are roughly 18 edges requiring one load, one store, and two reductions. The kernel should be efficient enough, though alternative approaches like vectorized loads could help, but the data isn't guaranteed contiguous across edges. The real bottleneck is the Python loop iterating over nodes sequentially. A CUDA implementation should provide significant speedup. Using 128 threads per block gives 3072 total threads across 24 blocks—lower occupancy but still exploiting parallelism. I could switch to 64 threads for 48 blocks to improve occupancy somewhat. The edge_scores tensor is float32 with shape (num_edges, num_heads), so for a given edge index, stride access across heads maintains memory locality. Each warp handles one (node, head) pair, loading edges sequentially as edge_scores[edge_idx, head] where edge_idx spans node_start to node_end. The stride is num_heads, yielding a regular access pattern. For edge indices within a node, consecutive edges map to consecutive memory locations (head 0 at offset 0, head 1 at offset 8, etc.), which the GPU can coalesce when multiple warps access different node-head pairs. The implementation looks correct, but there's a critical oversight: __expf has limited precision compared to std::exp. I should replace it with exp for better accuracy, though the performance difference is typically minimal. Now I need to construct the Python wrapper. The forward function receives row_ptr and edge_scores tensors, extracts the number of nodes and edges from their shapes, launches the CUDA kernel with appropriate block and grid dimensions, then returns the output tensor. The kernel calculates total warps needed based on nodes and heads, ensuring the grid size accounts for all warps per block. The softmax operation operates independently across edges, so the warp-level parallelism within each (node, head) pair handles the reduction and normalization correctly. Rather than using the warp ID directly, I should derive it from the thread index itself—computing `warp_id = idx / 32` and `lane = idx % 32`. This means `total_warps` must account for how warps distribute across blocks, where `total_warps = num_nodes * num_heads` and each block processes multiple warps concurrently. The block indexing becomes `block_id = warp_id / 4` for 128-thread blocks, with `local_warp = warp_id % 4` determining which warp within the block. For 64-thread blocks, the division shifts to `warp_id / 2` instead. Looking at the actual thread indices: with 128 threads per block, blockIdx 0 handles threads 0-127, where the first warp (threads 0-31) processes warp_id 0, the second warp (threads 32-63) processes warp_id 1, and so on through warp_id 3. With 64 threads per block, blockIdx 0 covers threads 0-63, with the first warp (threads 0-31) on warp_id 0 and the second warp (threads 32-63) on warp_id 1. For 3072 total warps across 8 heads, 128 threads yields 24 blocks while 64 threads gives 48 blocks—both viable configurations. There's also a synchronization consideration with __shfl_xor_sync, which handles warp-level operations. Since all threads in each warp are active during the reduction (no divergence within warps, only across them), the synchronization primitive should work correctly. The implementation handles empty nodes naturally: the start and end pointers will be equal, so num_edges becomes 0 and the loop never executes, leaving the output unmodified. Since the output tensor is created with torch.empty_like, it may contain garbage values, but the reference behavior for empty rows should be verified. In PyTorch, softmax on an empty tensor returns the tensor unchanged (empty output), which aligns with the kernel's behavior of simply not processing those indices. The main concern is whether degree=0 rows exist in the input, but the reference implementation generates degrees clamped to a minimum of 1, ensuring all rows have at least one element. Now I need to complete the final implementation. The CUDA kernel uses warp-level primitives to compute softmax across edges for each (node, head) pair. Each thread identifies which node and head it's processing, determines the edge range from the row pointer array, and then collaborates with other threads in its warp to find the maximum value, compute the sum of exponentials, and normalize the scores. score - max_val; local_sum += expf(exp_score); } start_e += 32; } float sum = WarpReduceMax(local_sum); sum = WarpReduceSum(local_sum); start_e = lane; while (start_e < num_edges) { int e = start_e + node_start; float score = edge_scores[e * num_heads + head]; out[e * num_heads + head] = expf(score - max_val) / sum; start_e += 32; } } I'm iterating through the edges assigned to this thread, computing the exponential of each score after subtracting the maximum for numerical stability, then summing these values. The sum gets reduced across the warp using shuffle operations, then I normalize each score by dividing by the total sum and write the result to the output array.