submission 128805

  #!POPCORN leaderboard vectorsum_v2
- #!POPCORN gpus H100
+ #!POPCORN gpus B200
   
  #!/usr/bin/env python3
  import torch
  ⋯ 14 unchanged lines
  #include <cuda_runtime.h>
  #include <ATen/ATen.h>
   
- // 1D reduction: sum(a[i]) into a scalar, with per-block partials
+ // 1D reduction: sum(a[i]) into a scalar using one atomicAdd per block
  template <typename scalar_t, typename acc_t>
  __global__ void vectorsum_kernel(const scalar_t* __restrict__ a,
-                                  acc_t* __restrict__ partial,
+                                  acc_t* __restrict__ out,
                                   long n) {
      extern __shared__ acc_t smem[];
   
-     const int tid     = threadIdx.x;
-     const int block_n = blockDim.x * gridDim.x;
-     long idx          = blockIdx.x * blockDim.x + tid;
+     const int tid         = threadIdx.x;
+     const int threads     = blockDim.x;
+     const int blocks      = gridDim.x;
+     const long grid_stride = static_cast<long>(threads) * static_cast<long>(blocks);
   
-     // 1) Per-thread strided accumulation in higher-precision acc_t
+     long idx = static_cast<long>(blockIdx.x) * threads + tid;
+  
+     // ---- CHANGED PART: batched grid-stride accumulation ----
+     constexpr int BATCH = 8;  // tune: 2, 4, 8, ...
+  
      acc_t local = static_cast<acc_t>(0);
-     for (; idx < n; idx += block_n) {
-         local += static_cast<acc_t>(a[idx]);
+     // We now stride in chunks of BATCH * grid_stride,
+     // and inside each iteration do BATCH loads.
+     for (long base = idx; base < n; base += static_cast<long>(grid_stride) * BATCH) {
+     #pragma unroll
+         for (int j = 0; j < BATCH; ++j) {
+             long i = base + static_cast<long>(j) * grid_stride;
+             if (i < n) {
+                 local += static_cast<acc_t>(a[i]);
+             }
+         }
      }
+     // ---- END CHANGED PART ----
   
      // 2) Write to shared memory
      smem[tid] = local;
      __syncthreads();
   
      // 3) Block-wide reduction in shared memory
-     for (int s = blockDim.x / 2; s > 0; s >>= 1) {
-         if (tid < s) {
-             smem[tid] += smem[tid + s];
+     for (int offset = threads / 2; offset > 0; offset >>= 1) {
+         if (tid < offset) {
+             smem[tid] += smem[tid + offset];
          }
          __syncthreads();
      }
   
-     // 4) Thread 0 writes the block's partial sum
+     // 4) Thread 0 atomically accumulates the block's sum into the global scalar
      if (tid == 0) {
-         partial[blockIdx.x] = smem[0];
+         atomicAdd(out, smem[0]);
      }
  }
   
- // Definition of vectorsum_cuda declared in CPP_SRC
+ // Public entrypoint called from Python
  torch::Tensor vectorsum_cuda(torch::Tensor a) {
+     TORCH_CHECK(a.is_cuda(), "vectorsum_cuda: input must be a CUDA tensor");
+     TORCH_CHECK(a.is_floating_point(), "vectorsum_cuda: input must be floating point");
+  
      auto a_c = a.contiguous();
      const long n = a_c.numel();
   
+     // Handle empty input: just return 0.0f on the same device
+     auto opts = torch::dtype(torch::kFloat32).device(a_c.device());
+     auto out  = torch::zeros({}, opts);  // scalar tensor
+  
      if (n == 0) {
-         // Handle empty input: return 0.0f on same device
-         auto empty_out = torch::zeros({}, a_c.options().dtype(torch::kFloat32));
-         return empty_out;
+         return out;
      }
   
-     constexpr int threads = 256;
+     using acc_t = float;  // accumulate in FP32
+  
+     const int threads = 256;
+     const int max_blocks = 1024;
      int blocks = static_cast<int>((n + threads - 1) / threads);
-     if (blocks < 1)   blocks = 1;
-     if (blocks > 1024) blocks = 1024;  // Safety cap
+     if (blocks > max_blocks) {
+         blocks = max_blocks;
+     }
   
-     // Partial sums kept in float64 for better numerical agreement with reference
-     auto partial = torch::empty({blocks}, a_c.options().dtype(torch::kFloat64));
+     const size_t shmem_bytes = threads * sizeof(acc_t);
   
-     AT_DISPATCH_FLOATING_TYPES_AND_HALF(a_c.scalar_type(), "vectorsum_cuda", [&] {
-         using acc_t = double;
+     AT_DISPATCH_FLOATING_TYPES_AND2(at::kHalf, at::kBFloat16,
+                                     a_c.scalar_type(), "vectorsum_cuda", [&] {
+         const scalar_t* a_ptr = a_c.data_ptr<scalar_t>();
+         acc_t* out_ptr = out.data_ptr<acc_t>();
+  
          vectorsum_kernel<scalar_t, acc_t>
-             <<<blocks, threads, threads * sizeof(acc_t)>>>(
-                 a_c.data_ptr<scalar_t>(),
-                 partial.data_ptr<acc_t>(),
+             <<<blocks, threads, shmem_bytes>>>(
+                 a_ptr,
+                 out_ptr,
                  n);
      });
   
-     // Final reduction on GPU in float64, then cast to float32
-     auto total64 = partial.sum();                       // float64 scalar, same device as a_c
-     auto total32 = total64.to(torch::kFloat32);         // match reference dtype
-  
-     return total32;
+     return out;
  }
  """
   
- # 3) Build extension once at import time.
+ # 3) Build the inline extension
  vectorsum_mod = load_inline(
-     name="vectorsum_ext",
-     cpp_sources=CPP_SRC,
-     cuda_sources=CUDA_SRC,
+     name="vectorsum_mod_batched",
+     cpp_sources=[CPP_SRC],
+     cuda_sources=[CUDA_SRC],
      functions=["vectorsum_cuda"],
-     with_cuda=True,
+     extra_cflags=["-O3"],
+     extra_cuda_cflags=["-O3"],
      verbose=False,
  )
   
-  
- @torch.inference_mode()
+ # 4) Popcorn entrypoint
  def custom_kernel(data):
      """
-     GPU Mode entrypoint.
+     Popcorn entrypoint.
   
-     The harness passes (input_tensor, output_tensor).
-  
-     We:
-       - grab the input tensor,
-       - ensure it's on CUDA,
-       - call the CUDA vectorsum kernel,
-       - return a single scalar tensor (float32 on CUDA).
+     data is a (x, out_buf) tuple:
+       - x: 1D CUDA tensor to reduce
+       - out_buf: preallocated output buffer (ignored here; we return our own scalar)
      """
-     x, out_buf = data  # out_buf is preallocated but we don't need to use it
+     x, out_buf = data
   
-     # Ensure tensor is on CUDA; generator already does this, but be robust
-     x = x.cuda(non_blocking=True)
+     if not x.is_cuda:
+         x = x.cuda(non_blocking=True)
   
      out = vectorsum_mod.vectorsum_cuda(x)
-  
-     # Could also do out_buf.copy_(out) and return out_buf, but returning `out` is fine
      return out