submission 772765

  #!POPCORN leaderboard vectorsum_v2
- #!POPCORN gpu A100
+ #!POPCORN gpu H100
   
  import torch
  from torch.utils.cpp_extension import load_inline
  from task import input_t, output_t
   
- # High-performance vector sum reduction using CUDA.
- # Strategy:
- #   1. Vectorized float4 loads to maximize memory bandwidth
- #   2. Thread coarsening: each thread processes multiple float4 chunks
- #   3. Warp-level shuffle reduction (no shared memory bank conflicts)
- #   4. Block-level reduction via shared memory
- #   5. Two-phase: phase 1 produces partial sums, phase 2 reduces them
+ # Vector sum reduction tuned for NVIDIA H100 (Hopper, sm_90)
  #
- # Tuned for A100 (108 SMs, 2 TB/s HBM2e bandwidth, 40MB L2).
+ # H100 SXM specs:
+ #   - 132 SMs, 128 CUDA cores/SM = 16896 cores
+ #   - ~3.35 TB/s HBM3 bandwidth
+ #   - 50 MB L2 cache
+ #   - Max 2048 threads/SM, 64 warps/SM
+ #
+ # Key tuning decisions vs A100:
+ #   - More blocks (528 = 132 SMs × 4) to saturate the larger GPU
+ #   - 256 threads/block, high occupancy
+ #   - H100's massive bandwidth means we're even more memory-bound
+ #   - Larger L2 cache helps with phase 2 partial sum reads
   
  cuda_source = r"""
  #include <cuda_runtime.h>
- #include <cuda_fp16.h>
   
- // Warp-level reduction using shuffle
  __device__ __forceinline__ float warp_reduce_sum(float val) {
      #pragma unroll
      for (int offset = 16; offset > 0; offset >>= 1) {
  ⋯ 2 unchanged lines
      return val;
  }
   
- // Phase 1: Each block reduces a large chunk of the input into one partial sum.
- // Uses float4 vectorized loads and thread coarsening.
+ __global__ void reduce_tiny(
+     const float* __restrict__ input,
+     float* __restrict__ output,
+     int n
+ ) {
+     __shared__ float smem[32];
+     const int tid = threadIdx.x;
+     float sum = 0.0f;
+  
+     for (int i = tid; i < n; i += blockDim.x) {
+         sum += input[i];
+     }
+  
+     sum = warp_reduce_sum(sum);
+     const int lane = tid & 31;
+     const int warp_id = tid >> 5;
+  
+     if (lane == 0) smem[warp_id] = sum;
+     __syncthreads();
+  
+     const int num_warps = (blockDim.x + 31) / 32;
+     if (warp_id == 0) {
+         sum = (lane < num_warps) ? smem[lane] : 0.0f;
+         sum = warp_reduce_sum(sum);
+     }
+     if (tid == 0) output[0] = sum;
+ }
+  
  __global__ void reduce_phase1(
      const float* __restrict__ input,
      float* __restrict__ partial_sums,
      int n
  ) {
-     // Shared memory for block-level reduction (one slot per warp)
-     __shared__ float smem[32];  // max 32 warps per block
-  
+     __shared__ float smem[32];
      const int tid = threadIdx.x;
      const int bid = blockIdx.x;
      const int block_size = blockDim.x;
  ⋯ 1 unchanged lines
   
      float sum = 0.0f;
   
-     // Vectorized float4 loads: each thread processes 4 floats at a time
-     // with grid-stride loop for coarsening
-     const int n4 = n / 4;
+     // float4 vectorized loads with grid-stride loop
+     const int n4 = n >> 2;
      const float4* input4 = reinterpret_cast<const float4*>(input);
-  
      for (int i = bid * block_size + tid; i < n4; i += grid_size) {
          float4 v = input4[i];
          sum += v.x + v.y + v.z + v.w;
      }
   
-     // Handle remaining elements (n % 4 tail)
-     int tail_start = n4 * 4;
+     // Tail
+     int tail_start = n4 << 2;
      for (int i = tail_start + bid * block_size + tid; i < n; i += grid_size) {
          sum += input[i];
      }
   
-     // Warp-level reduction
      sum = warp_reduce_sum(sum);
-  
-     // Write warp results to shared memory
      const int lane = tid & 31;
      const int warp_id = tid >> 5;
   
-     if (lane == 0) {
-         smem[warp_id] = sum;
-     }
+     if (lane == 0) smem[warp_id] = sum;
      __syncthreads();
   
-     // First warp reduces all warp partial sums
      const int num_warps = (block_size + 31) / 32;
      if (warp_id == 0) {
          sum = (lane < num_warps) ? smem[lane] : 0.0f;
          sum = warp_reduce_sum(sum);
      }
-  
-     // Thread 0 writes the block's partial sum
-     if (tid == 0) {
-         partial_sums[bid] = sum;
-     }
+     if (tid == 0) partial_sums[bid] = sum;
  }
   
- // Phase 2: Reduce partial sums into a single scalar.
- // Single block, single warp is enough for small number of partial sums.
  __global__ void reduce_phase2(
      const float* __restrict__ partial_sums,
      float* __restrict__ output,
      int n
  ) {
      __shared__ float smem[32];
-  
      const int tid = threadIdx.x;
      float sum = 0.0f;
   
  ⋯ 2 unchanged lines
      }
   
      sum = warp_reduce_sum(sum);
-  
      const int lane = tid & 31;
      const int warp_id = tid >> 5;
   
-     if (lane == 0) {
-         smem[warp_id] = sum;
-     }
+     if (lane == 0) smem[warp_id] = sum;
      __syncthreads();
   
      const int num_warps = (blockDim.x + 31) / 32;
  ⋯ 1 unchanged lines
          sum = (lane < num_warps) ? smem[lane] : 0.0f;
          sum = warp_reduce_sum(sum);
      }
-  
-     if (tid == 0) {
-         output[0] = sum;
-     }
+     if (tid == 0) output[0] = sum;
  }
   
  torch::Tensor vector_sum_cuda(torch::Tensor input, torch::Tensor output) {
      const int n = input.numel();
   
-     // A100: 108 SMs. We want enough blocks to saturate all SMs
-     // but not so many that phase 2 becomes expensive.
-     // 256 threads/block, ~432 blocks gives good occupancy on A100.
+     if (n <= 1024) {
+         reduce_tiny<<<1, 256>>>(
+             input.data_ptr<float>(), output.data_ptr<float>(), n);
+         return output;
+     }
+  
+     // H100: 132 SMs, 256 threads/block
      const int threads = 256;
-     const int blocks = min(432, (n / 4 + threads - 1) / threads);
+     int blocks;
   
-     // Allocate temporary buffer for partial sums
+     if (n <= 32768) {
+         blocks = min(132, (n / 4 + threads - 1) / threads);
+     } else {
+         // 132 SMs × 4 blocks/SM = 528 blocks for full occupancy
+         blocks = min(528, (n / 4 + threads - 1) / threads);
+     }
+  
      auto partial_sums = torch::empty({blocks}, input.options());
   
      reduce_phase1<<<blocks, threads>>>(
-         input.data_ptr<float>(),
-         partial_sums.data_ptr<float>(),
-         n
-     );
+         input.data_ptr<float>(), partial_sums.data_ptr<float>(), n);
   
-     // Phase 2: reduce partial sums (blocks <= 432, one block of 256 threads is plenty)
-     reduce_phase2<<<1, 256>>>(
-         partial_sums.data_ptr<float>(),
-         output.data_ptr<float>(),
-         blocks
-     );
+     int p2_threads = min(256, ((blocks + 31) / 32) * 32);
+     reduce_phase2<<<1, p2_threads>>>(
+         partial_sums.data_ptr<float>(), output.data_ptr<float>(), blocks);
   
      return output;
  }
  ⋯ 4 unchanged lines
  torch::Tensor vector_sum_cuda(torch::Tensor input, torch::Tensor output);
  """
   
- # Compile with optimizations
  module = load_inline(
-     name="vector_sum_fast",
+     name="vector_sum_h100",
      cpp_sources=cpp_source,
      cuda_sources=cuda_source,
      functions=["vector_sum_cuda"],