submission 771392

  ⋯ 4 unchanged lines
  import triton
  import triton.language as tl
   
+  
  @triton.jit
- def sum_kernel_a100(
-     x_ptr,
+ def reduce_chunked_kernel(
+     in_ptr,
      out_ptr,
      n_elements,
      BLOCK_SIZE: tl.constexpr,
+     CHUNKS_PER_PROGRAM: tl.constexpr,
  ):
-     """
-     使用共享内存进行块内归约，并将最终结果原子累加到全局输出。
-     所有归约均在 float64 下完成，确保与参考实现数值一致。
-     """
      pid = tl.program_id(axis=0)
-     block_start = pid * BLOCK_SIZE
-     offsets = block_start + tl.arange(0, BLOCK_SIZE)
-     mask = offsets < n_elements
+     base = pid * BLOCK_SIZE * CHUNKS_PER_PROGRAM
+     lane = tl.arange(0, BLOCK_SIZE)
   
-     x = tl.load(x_ptr + offsets, mask=mask, other=0.0)
+     acc = tl.zeros((BLOCK_SIZE,), dtype=tl.float32)
+     for i in range(CHUNKS_PER_PROGRAM):
+         offsets = base + i * BLOCK_SIZE + lane
+         mask = offsets < n_elements
+         x = tl.load(in_ptr + offsets, mask=mask, other=0.0)
+         acc += x
   
-     # 块内归约：使用 tl.sum 在寄存器层完成（Triton 会自动优化）
-     partial_sum = tl.sum(x)
+     partial = tl.sum(acc, axis=0)
+     tl.store(out_ptr + pid, partial)
   
-     # 将本块的部分和原子累加到全局双精度输出
-     tl.atomic_add(out_ptr, partial_sum)
   
-  
- def custom_kernel(data: input_t) -> output_t:
+ def ref_kernel(data: input_t) -> output_t:
      """
-     A100 优化的全局求和，严格匹配 ref_kernel 的数值精度。
+     Reference implementation of vector sum reduction using PyTorch.
+     Args:
+         data: Input tensor to be reduced
+     Returns:
+         Tensor containing the sum of all elements
      """
+     with DeterministicContext():
+         data, output = data
+         # Let's be on the safe side here, and do the reduction in 64 bit
+         output = data.to(torch.float64).sum().to(torch.float32)
+         return output
+  
+  
+ def custom_kernel(data: input_t) -> output_t:
      input_tensor, _ = data
      n_elements = input_tensor.numel()
   
      if n_elements == 0:
          return torch.zeros((), device=input_tensor.device, dtype=torch.float32)
   
-     # 确保输入连续
      x = input_tensor.contiguous()
   
-     # 双精度临时缓冲区，用于原子累加
-     out_double = torch.zeros(1, device="cuda", dtype=torch.float64)
+     # A100 优化参数：每个 program 处理 4 个 1024 元素块（共 4096 元素）
+     # 目标是减少 launch program 数，同时保持较高并行度和内存吞吐。
+     BLOCK_SIZE = 1024
+     CHUNKS_PER_PROGRAM = 2
+     num_warps = 8
+     num_stages = 4
   
-     # A100 调优参数
-     BLOCK_SIZE = 2048
-     grid = lambda meta: (triton.cdiv(n_elements, meta["BLOCK_SIZE"]),)
-  
-     # 正确传递 num_warps 和 num_stages（作为 kernel 启动参数）
-     sum_kernel_a100[grid](
-         x, out_double, n_elements,
+     n_partials = triton.cdiv(n_elements, BLOCK_SIZE * CHUNKS_PER_PROGRAM)
+     partials = torch.empty((n_partials,), device=x.device, dtype=torch.float32)
+     reduce_chunked_kernel[(n_partials,)](
+         x,
+         partials,
+         n_elements,
          BLOCK_SIZE=BLOCK_SIZE,
-         num_warps=8,      # A100 每个 SM 最多 8 warps
-         num_stages=4,     # 软件流水线深度
+         CHUNKS_PER_PROGRAM=CHUNKS_PER_PROGRAM,
+         num_warps=num_warps,
+         num_stages=num_stages,
      )
+     input = partials
+     output = x
   
-     return out_double[0].to(torch.float32)
+     while n_partials > 1:
+         next_n = triton.cdiv(n_partials, BLOCK_SIZE * CHUNKS_PER_PROGRAM)
+         #next_partials = torch.empty((next_n,), device=x.device, dtype=torch.float32)
+         reduce_chunked_kernel[(next_n,)](
+             input,
+             output,  # 重用输入缓冲区作为中间结果存储
+             n_partials,
+             BLOCK_SIZE=BLOCK_SIZE,
+             CHUNKS_PER_PROGRAM=CHUNKS_PER_PROGRAM,
+             num_warps=num_warps,
+             num_stages=num_stages,
+         )
+         # 交换输入输出缓冲区
+         input, output = output, input
+         n_partials = next_n
   
+     return input[0].to(torch.float32)
   
- def ref_kernel(data: input_t) -> output_t:
-     """
-     Reference implementation of vector sum reduction using PyTorch.
-     Args:
-         data: Input tensor to be reduced
-     Returns:
-         Tensor containing the sum of all elements
-     """
-     with DeterministicContext():
-         data, output = data
-         # Let's be on the safe side here, and do the reduction in 64 bit
-         output = data.to(torch.float64).sum().to(torch.float32)
-         return output
   
  def generate_input(size: int, seed: int) -> input_t:
      """
  ⋯ 30 unchanged lines
   
   
  check_implementation = make_match_reference(ref_kernel)
-