submission 757709

  ⋯ 4 unchanged lines
  import triton.language as tl
  from typing import Tuple
   
- # ----------------------------------------------------------------------
- # Combine function for associative scan (simple addition)
- # ----------------------------------------------------------------------
+ # ------------------------------------------------------------------
+ # Kernel: compute per‑block sums
+ # ------------------------------------------------------------------
  @triton.jit
- def _add_combine(a, b):
-     return a + b
-  
- # ----------------------------------------------------------------------
- # Kernel 1: compute per‑block sums of the input tensor
- # ----------------------------------------------------------------------
- @triton.jit
  def block_sum_kernel(
-     x_ptr,                # *Pointer* to input tensor
-     block_sum_ptr,        # *Pointer* to per‑block sum buffer
-     n: tl.int64,          # total number of elements
+     x_ptr,               # *float32, input vector
+     block_sum_ptr,       # *float32, per‑block sums (output)
+     N,                   # int32, total number of elements
      BLOCK_SIZE: tl.constexpr,
  ):
      pid = tl.program_id(0)                     # block index
      block_start = pid * BLOCK_SIZE
      offsets = block_start + tl.arange(0, BLOCK_SIZE)
-     mask = offsets < n
+     mask = offsets < N
   
-     # Load a chunk of the input (zero‑pad out‑of‑bounds)
+     # Load a tile (zero‑pad out‑of‑bounds)
      x = tl.load(x_ptr + offsets, mask=mask, other=0.0)
   
-     # Reduce to a single sum for this block
+     # Reduce across the block
      block_sum = tl.sum(x)
   
-     # Write the per‑block sum
+     # Store the block sum
      tl.store(block_sum_ptr + pid, block_sum)
   
- # ----------------------------------------------------------------------
- # Kernel 2: final inclusive scan + add exclusive per‑block offset
- # ----------------------------------------------------------------------
+  
+ # ------------------------------------------------------------------
+ # Kernel: final inclusive scan using per‑block exclusive offsets
+ # ------------------------------------------------------------------
  @triton.jit
  def final_scan_kernel(
-     x_ptr,                # *Pointer* to input tensor
-     out_ptr,              # *Pointer* to output tensor
-     block_sum_ptr,        # *Pointer* to per‑block sum buffer (original)
-     block_sum_scan_ptr,   # *Pointer* to inclusive scan of per‑block sums
-     n: tl.int64,          # total number of elements
+     x_ptr,               # *float32, input vector
+     out_ptr,             # *float32, output vector (final result)
+     block_offsets_ptr,   # *float32, inclusive per‑block offsets
+     block_sums_ptr,      # *float32, per‑block sums
+     N,                   # int32, total number of elements
      BLOCK_SIZE: tl.constexpr,
  ):
      pid = tl.program_id(0)                     # block index
      block_start = pid * BLOCK_SIZE
      offsets = block_start + tl.arange(0, BLOCK_SIZE)
-     mask = offsets < n
+     mask = offsets < N
   
-     # Load the input slice for this block
+     # Load input tile
      x = tl.load(x_ptr + offsets, mask=mask, other=0.0)
   
-     # Inclusive scan within the block
-     local_scan = tl.associative_scan(x, axis=0, combine_fn=_add_combine)
+     # Inclusive prefix sum inside the block
+     y = tl.cumsum(x)
   
-     # Load per‑block sums (original) and its inclusive scan
-     block_sum = tl.load(block_sum_ptr + pid)
-     block_sum_scan = tl.load(block_sum_scan_ptr + pid)
+     # Load inclusive offset for this block and its block sum
+     offset_incl = tl.load(block_offsets_ptr + pid)
+     block_sum = tl.load(block_sums_ptr + pid)
   
-     # Exclusive offset = sum of all previous blocks
-     offset = block_sum_scan - block_sum
+     # Convert inclusive offset to exclusive offset
+     offset_excl = offset_incl - block_sum
   
-     # Apply the offset to the locally scanned values
-     out = local_scan + offset
+     # Apply the exclusive offset
+     y = y + offset_excl
   
-     # Write the final result
-     tl.store(out_ptr + offsets, out, mask=mask)
+     # Write results
+     tl.store(out_ptr + offsets, y, mask=mask)
   
- # ----------------------------------------------------------------------
- # Public API
- # ----------------------------------------------------------------------
+  
+ # ------------------------------------------------------------------
+ # Kernel: single‑pass scan for very small inputs (N <= BLOCK_SIZE)
+ # ------------------------------------------------------------------
+ @triton.jit
+ def small_scan_kernel(
+     x_ptr,               # *float32, input vector
+     out_ptr,             # *float32, output vector
+     N,                   # int32, total number of elements
+     BLOCK_SIZE: tl.constexpr,
+ ):
+     offsets = tl.arange(0, BLOCK_SIZE)
+     mask = offsets < N
+     x = tl.load(x_ptr + offsets, mask=mask, other=0.0)
+     y = tl.cumsum(x)
+     tl.store(out_ptr + offsets, y, mask=mask)
+  
+  
+ # ------------------------------------------------------------------
+ # Public interface
+ # ------------------------------------------------------------------
  def custom_kernel(data: Tuple[torch.Tensor, torch.Tensor]) -> torch.Tensor:
      """
-     Inclusive prefix sum (cumsum) over a 1‑D float32 tensor using a two‑pass
-     Triton implementation with a small‑size fallback to torch.cumsum.
+     Inclusive prefix sum (scan) over a 1‑D float32 tensor using Triton.
      """
      x, output = data
   
-     # ------------------------------------------------------------------
-     # Validation
-     # ------------------------------------------------------------------
-     if not (x.is_cuda and output.is_cuda):
-         raise RuntimeError("custom_kernel expects CUDA tensors")
-     if x.dtype != torch.float32 or output.dtype != torch.float32:
-         raise RuntimeError("custom_kernel only supports float32 tensors")
-     if x.shape != output.shape:
-         raise ValueError("input and output must have the same shape")
-     if x.ndim != 1:
-         raise ValueError("custom_kernel only supports 1‑D tensors")
+     # Basic sanity checks
+     assert x.is_contiguous() and output.is_contiguous(), "Both tensors must be contiguous"
+     assert x.dtype == torch.float32 and output.dtype == torch.float32, "Only float32 supported"
+     assert x.shape == output.shape, "Input and output must have the same shape"
   
-     n = x.numel()
-     if n == 0:
+     N = x.shape[0]
+     if N == 0:
+         output.zero_()
          return output
   
      # ------------------------------------------------------------------
-     # Small‑tensor fast path using native torch.cumsum
+     # Configuration (tuned for Hopper/H200)
      # ------------------------------------------------------------------
-     SMALL_THRESHOLD = 1 << 20  # 1,048,576 elements
-     if n <= SMALL_THRESHOLD:
-         torch.cumsum(x, dim=0, out=output)
-         return output
+     BLOCK_SIZE = 1024                     # max threads per block
+     num_blocks = (N + BLOCK_SIZE - 1) // BLOCK_SIZE
   
      # ------------------------------------------------------------------
-     # Triton launch configuration
+     # Small‑input fast path
      # ------------------------------------------------------------------
-     BLOCK_SIZE = 32768                 # elements per block (multiple of 32)
-     NUM_WARPS_SUM = 8                  # 8 warps → 256 threads per block (reduction)
-     NUM_WARPS_SCAN = 8                 # 8 warps → 256 threads per block (scan)
-     num_blocks = (n + BLOCK_SIZE - 1) // BLOCK_SIZE
+     if N <= BLOCK_SIZE:
+         small_scan_kernel[(1,)](
+             x, output, N,
+             BLOCK_SIZE=BLOCK_SIZE,
+             num_stages=2,
+         )
+         return output
   
-     # Buffer for per‑block sums
-     block_sums = torch.empty(num_blocks, dtype=x.dtype, device=x.device)
-  
      # ------------------------------------------------------------------
-     # Pass 1: compute per‑block sums
+     # Large‑input hierarchical scan
      # ------------------------------------------------------------------
+     # 1️⃣ Compute per‑block sums
+     block_sums = torch.empty(num_blocks, dtype=x.dtype, device=x.device)
      block_sum_kernel[(num_blocks,)](
-         x,
-         block_sums,
-         n,
+         x, block_sums, N,
          BLOCK_SIZE=BLOCK_SIZE,
-         num_warps=NUM_WARPS_SUM,
-         num_stages=3,          # modest pipelining for the reduction kernel
+         num_stages=2,
      )
   
-     # ------------------------------------------------------------------
-     # Scan the per‑block totals on the GPU (inclusive)
-     # ------------------------------------------------------------------
-     block_sums_scan = torch.cumsum(block_sums, dim=0)
+     # 2️⃣ Compute inclusive block offsets on the GPU
+     block_offsets = torch.empty_like(block_sums)
+     torch.cumsum(block_sums, dim=0, out=block_offsets)
   
-     # ------------------------------------------------------------------
-     # Pass 2: local inclusive scan + add exclusive per‑block offsets
-     # ------------------------------------------------------------------
+     # 3️⃣ Final scan with exclusive offsets
      final_scan_kernel[(num_blocks,)](
-         x,
-         output,
-         block_sums,
-         block_sums_scan,
-         n,
+         x, output, block_offsets, block_sums, N,
          BLOCK_SIZE=BLOCK_SIZE,
-         num_warps=NUM_WARPS_SCAN,
-         num_stages=6,          # deeper pipelining to hide memory latency
+         num_stages=4,
      )
-  
      return output
  # EVOLVE-BLOCK-END
   No newline at end of file