submission 490611

+ import torch
  import triton
  import triton.language as tl
- import torch
   
+  
+ @triton.autotune(
+     configs=[
+         triton.Config({"BLOCK_M": 64, "BLOCK_N": 64, "BLOCK_K": 32, "GROUP_SIZE_M": 8}, num_stages=3, num_warps=4),
+         triton.Config({"BLOCK_M": 64, "BLOCK_N": 128, "BLOCK_K": 32, "GROUP_SIZE_M": 8}, num_stages=3, num_warps=4),
+         triton.Config({"BLOCK_M": 128, "BLOCK_N": 64, "BLOCK_K": 32, "GROUP_SIZE_M": 8}, num_stages=3, num_warps=4),
+         triton.Config({"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 32, "GROUP_SIZE_M": 8}, num_stages=4, num_warps=8),
+         triton.Config({"BLOCK_M": 32, "BLOCK_N": 128, "BLOCK_K": 32, "GROUP_SIZE_M": 8}, num_stages=3, num_warps=4),
+         triton.Config({"BLOCK_M": 64, "BLOCK_N": 256, "BLOCK_K": 32, "GROUP_SIZE_M": 8}, num_stages=4, num_warps=8),
+     ],
+     key=["M", "N", "K"],
+ )
  @triton.jit
  def _matmul_kernel(
      a_ptr, b_ptr, c_ptr,
  ⋯ 1 unchanged lines
      stride_am, stride_ak,
      stride_bk, stride_bn,
      stride_cm, stride_cn,
-     BLOCK_SIZE_M: tl.constexpr,
-     BLOCK_SIZE_N: tl.constexpr,
-     BLOCK_SIZE_K: tl.constexpr,
+     BLOCK_M: tl.constexpr,
+     BLOCK_N: tl.constexpr,
+     BLOCK_K: tl.constexpr,
+     GROUP_SIZE_M: tl.constexpr,
  ):
-     """
-     Triton kernel for matrix multiplication C = A @ B
-     
-     A is (M, K), B is (K, N), C is (M, N)
-     Each program computes a BLOCK_SIZE_M x BLOCK_SIZE_N tile of C.
-     """
-     # Get program IDs
-     pid_m = tl.program_id(0)
-     pid_n = tl.program_id(1)
-     
-     # Calculate the starting row and column for this block
-     rm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
-     rn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
-     
-     # Initialize accumulator in float32 for numerical stability
-     acc = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
-     
-     # Iterate over K dimension in blocks
-     for k_start in range(0, K, BLOCK_SIZE_K):
-         rk = k_start + tl.arange(0, BLOCK_SIZE_K)
-         
-         # Load block of A: shape (BLOCK_SIZE_M, BLOCK_SIZE_K)
-         # A[rm, rk] where rm is rows, rk is columns
-         a_ptrs = a_ptr + rm[:, None] * stride_am + rk[None, :] * stride_ak
-         a_mask = (rm[:, None] < M) & (rk[None, :] < K)
+     # 2D tiling with grouping along M to improve L2 locality. Single-axis launch.
+     pid = tl.program_id(axis=0)
+     num_pid_m = tl.cdiv(M, BLOCK_M)
+     num_pid_n = tl.cdiv(N, BLOCK_N)
+     num_pid_in_group = GROUP_SIZE_M * num_pid_n
+  
+     group_id = pid // num_pid_in_group
+     first_pid_m = group_id * GROUP_SIZE_M
+     group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)
+     pid_m = first_pid_m + (pid % group_size_m)
+     pid_n = (pid % num_pid_in_group) // group_size_m
+  
+     # Compute tile offsets
+     offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
+     offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
+     offs_k = tl.arange(0, BLOCK_K)
+  
+     # Help codegen with alignment/contiguity hints
+     offs_m = tl.max_contiguous(tl.multiple_of(offs_m, BLOCK_M), BLOCK_M)
+     offs_n = tl.max_contiguous(tl.multiple_of(offs_n, BLOCK_N), BLOCK_N)
+  
+     # Accumulator in FP32
+     acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+  
+     # Loop over K tiles
+     k_tiles = tl.cdiv(K, BLOCK_K)
+     for kt in range(0, k_tiles):
+         k_start = kt * BLOCK_K
+         # Pointers to A and B tiles
+         a_ptrs = a_ptr + (offs_m[:, None] * stride_am + (k_start + offs_k[None, :]) * stride_ak)
+         b_ptrs = b_ptr + ((k_start + offs_k[:, None]) * stride_bk + offs_n[None, :] * stride_bn)
+  
+         # Masks for out-of-bounds on K/M/N
+         a_mask = (offs_m[:, None] < M) & ((k_start + offs_k[None, :]) < K)
+         b_mask = ((k_start + offs_k[:, None]) < K) & (offs_n[None, :] < N)
+  
+         # Load tiles from global memory
          a = tl.load(a_ptrs, mask=a_mask, other=0.0)
-         
-         # Load block of B: shape (BLOCK_SIZE_K, BLOCK_SIZE_N)
-         # B[rk, rn] where rk is rows, rn is columns
-         b_ptrs = b_ptr + rk[:, None] * stride_bk + rn[None, :] * stride_bn
-         b_mask = (rk[:, None] < K) & (rn[None, :] < N)
          b = tl.load(b_ptrs, mask=b_mask, other=0.0)
-         
-         # Accumulate: acc += A @ B for this K block
+  
+         # Multiply-accumulate on tensor cores
          acc = tl.dot(a, b, acc)
-     
-     # Convert accumulator back to output dtype and store
-     c = acc.to(tl.float16)
-     
-     # Store the result block
-     c_ptrs = c_ptr + rm[:, None] * stride_cm + rn[None, :] * stride_cn
-     c_mask = (rm[:, None] < M) & (rn[None, :] < N)
-     tl.store(c_ptrs, c, mask=c_mask)
   
+     # Write result tile to C
+     c_ptrs = c_ptr + (offs_m[:, None] * stride_cm + offs_n[None, :] * stride_cn)
+     c_mask = (offs_m[:, None] < M) & (offs_n[None, :] < N)
+     tl.store(c_ptrs, acc.to(c_ptr.dtype.element_ty), mask=c_mask)
   
- def kernel_function(data):
+  
+ def kernel_function(*args):
      """
-     Wrapper function for matrix multiplication.
-     
-     Takes a tuple (a, b, c) where:
-     - a: input matrix of shape (M, K), dtype float16
-     - b: input matrix of shape (K, N), dtype float16
-     - c: output matrix of shape (M, N), dtype float16 (pre-allocated but we overwrite)
-     
-     Computes C = A @ B using a fused Triton matmul kernel.
-     
-     Fusion note: This is a single matmul operation, so no additional fusion is needed.
-     The kernel fuses the entire M*K*N computation into a single kernel launch with
-     tiled accumulation for efficiency.
+     Triton matmul kernel wrapper.
+  
+     Accepts either:
+     - a single tuple/list: (a, b, c)
+     - three separate arguments: a, b, c
+  
+     Performs c = a @ b where:
+       - a: [M, K], float16 on CUDA
+       - b: [K, N], float16 on CUDA
+       - c: [M, N], float16 on CUDA (output buffer)
+  
+     Fusion reasoning:
+     - The test requires only a matrix multiplication. No bias or activation tensors are provided.
+     - We implement a single fused kernel that performs blocked loads, FP32 accumulation using tl.dot,
+       and a final store to FP16. There are no additional operator stages to fuse here.
+       Should future requirements include bias/activation, they can be fused in the epilogue.
+  
+     Runtime constraints:
+     - This wrapper only validates inputs, allocates/uses the output buffer, sets up strides,
+       computes the launch grid, and dispatches the Triton kernel. All math runs inside the kernel.
      """
-     a, b, c = data
-     
-     # Get dimensions
+     # Support both tuple-style and arg-style calls.
+     if len(args) == 1 and isinstance(args[0], (tuple, list)):
+         if len(args[0]) != 3:
+             raise TypeError("Expected a tuple/list of three tensors (a, b, c).")
+         a, b, c = args[0]
+     elif len(args) == 3:
+         a, b, c = args
+     else:
+         raise TypeError("kernel_function expects either (a, b, c) or a single tuple/list (a, b, c).")
+  
+     # Basic validation
+     if not isinstance(a, torch.Tensor) or not isinstance(b, torch.Tensor) or not isinstance(c, torch.Tensor):
+         raise TypeError("All inputs must be torch.Tensor instances.")
+     if a.device.type != "cuda" or b.device.type != "cuda" or c.device.type != "cuda":
+         raise RuntimeError("All tensors must be on CUDA device.")
+     if a.dtype != torch.float16 or b.dtype != torch.float16 or c.dtype != torch.float16:
+         raise RuntimeError("This kernel expects float16 tensors.")
+     if a.shape[1] != b.shape[0]:
+         raise RuntimeError(f"Incompatible shapes: a.shape={a.shape}, b.shape={b.shape} (K mismatch).")
+     if a.shape[0] != c.shape[0] or b.shape[1] != c.shape[1]:
+         raise RuntimeError(f"Output shape mismatch: c.shape={c.shape}, expected {(a.shape[0], b.shape[1])}.")
+  
+     # Shapes and strides
      M, K = a.shape
-     K2, N = b.shape
-     assert K == K2, f"Incompatible dimensions: a.shape[1]={K} != b.shape[0]={K2}"
-     
-     # Allocate output (use the provided c or create new one)
-     output = torch.empty((M, N), device=a.device, dtype=a.dtype)
-     
-     # Block sizes - chosen as powers of 2 for efficiency
-     # Using 32x32 blocks which work well for various sizes
-     BLOCK_SIZE_M = 32
-     BLOCK_SIZE_N = 32
-     BLOCK_SIZE_K = 32
-     
-     # Calculate grid dimensions
-     grid = (triton.cdiv(M, BLOCK_SIZE_M), triton.cdiv(N, BLOCK_SIZE_N))
-     
-     # Launch kernel
+     Kb, N = b.shape
+     assert K == Kb
+     # Allocate c if needed (test passes an existing buffer; we still guard it)
+     if c.numel() != M * N or c.shape != (M, N):
+         c = torch.empty((M, N), device=a.device, dtype=a.dtype)
+  
+     # Grid computation: 1D launch, flattened (pid_m, pid_n)
+     def grid(meta):
+         return (triton.cdiv(M, meta["BLOCK_M"]) * triton.cdiv(N, meta["BLOCK_N"]),)
+  
      _matmul_kernel[grid](
-         a, b, output,
+         a, b, c,
          M, N, K,
          a.stride(0), a.stride(1),
          b.stride(0), b.stride(1),
-         output.stride(0), output.stride(1),
-         BLOCK_SIZE_M=BLOCK_SIZE_M,
-         BLOCK_SIZE_N=BLOCK_SIZE_N,
-         BLOCK_SIZE_K=BLOCK_SIZE_K,
+         c.stride(0), c.stride(1),
      )
-     
-     return output
   
+     # Return the output tensor
+     return c
+  
  import inspect
   
  def custom_kernel(input):