submission 511441

  ⋯ 4 unchanged lines
   
  @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),
+         triton.Config({"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 32}, num_stages=3, num_warps=8),
+         triton.Config({"BLOCK_M": 64, "BLOCK_N": 64, "BLOCK_K": 32}, num_stages=3, num_warps=4),
+         triton.Config({"BLOCK_M": 64, "BLOCK_N": 128, "BLOCK_K": 32}, num_stages=3, num_warps=4),
+         triton.Config({"BLOCK_M": 128, "BLOCK_N": 64, "BLOCK_K": 32}, num_stages=3, num_warps=4),
+         triton.Config({"BLOCK_M": 32, "BLOCK_N": 128, "BLOCK_K": 32}, num_stages=2, num_warps=4),
+         triton.Config({"BLOCK_M": 64, "BLOCK_N": 256, "BLOCK_K": 32}, num_stages=4, num_warps=8),
      ],
      key=["M", "N", "K"],
  )
  ⋯ 7 unchanged lines
      BLOCK_M: tl.constexpr,
      BLOCK_N: tl.constexpr,
      BLOCK_K: tl.constexpr,
-     GROUP_SIZE_M: tl.constexpr,
  ):
-     # 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
+     """
+     Blocked matrix multiplication C = A @ B, with:
+       - A: [M, K]
+       - B: [K, N]
+       - C: [M, N]
+     Accumulation is performed in fp32 for numerical stability and stored to output dtype.
+     """
+     pid_m = tl.program_id(axis=0)
+     pid_n = tl.program_id(axis=1)
   
-     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
+     # Compute offsets for this program (block)
      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
+     # Iterate over K dimension
      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
+     for k_tile in range(0, k_tiles):
+         offs_k = k_tile * BLOCK_K + tl.arange(0, BLOCK_K)
+         # Pointers for the current tile of A and B
+         a_ptrs = a_ptr + (offs_m[:, None] * stride_am + offs_k[None, :] * stride_ak)
+         b_ptrs = b_ptr + (offs_k[:, None] * stride_bk + offs_n[None, :] * stride_bn)
+         # Masks for boundary conditions
+         a_mask = (offs_m[:, None] < M) & (offs_k[None, :] < K)
+         b_mask = (offs_n[None, :] < N) & (offs_k[:, None] < K)
+         # Load tiles; missing elements are set to 0
          a = tl.load(a_ptrs, mask=a_mask, other=0.0)
          b = tl.load(b_ptrs, mask=b_mask, other=0.0)
-  
-         # Multiply-accumulate on tensor cores
+         # Accumulate
          acc = tl.dot(a, b, acc)
   
-     # Write result tile to C
+     # Write back to C
+     c = acc.to(c_ptr.dtype.element_ty)
      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)
+     tl.store(c_ptrs, c, mask=c_mask)
   
   
- def kernel_function(*args):
+ def _extract_abc_from_args(*args):
      """
-     Triton matmul kernel wrapper.
+     Helper to support multiple calling conventions:
+       - kernel_function((a, b, c), c)
+       - kernel_function((a, b), c)
+       - kernel_function(a, b, c)
+       - kernel_function((a, b, c),)
+       - kernel_function(a, b)
+       - kernel_function((a, b),)
+     Returns (a, b, c or None).
+     """
+     a = b = c = None
+     if len(args) == 1 and isinstance(args[0], tuple):
+         data = args[0]
+         if len(data) == 3:
+             a, b, c = data
+         elif len(data) == 2:
+             a, b = data
+         else:
+             raise TypeError("Single tuple argument must be length 2 or 3.")
+     elif len(args) == 2:
+         if isinstance(args[0], tuple):
+             data = args[0]
+             if len(data) == 3:
+                 a, b, c_tuple = data
+                 c = args[1] if isinstance(args[1], torch.Tensor) else c_tuple
+             elif len(data) == 2:
+                 a, b = data
+                 c = args[1] if isinstance(args[1], torch.Tensor) else None
+             else:
+                 raise TypeError("Tuple must be length 2 or 3.")
+         else:
+             a, b = args
+             c = None
+     elif len(args) == 3:
+         a, b, c = args
+     else:
+         raise TypeError("Unsupported argument pattern for kernel_function.")
+     return a, b, c
   
-     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)
+ def kernel_function(*args):
+     """
+     Wrapper for a fused matmul kernel C = A @ B.
   
-     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.
+     Fusion discussion:
+     - The test requires matching plain matmul outputs for fp16 inputs and does not provide bias or activation tensors.
+     - As such, the only meaningful stage is the matmul itself; no additional operators are fused.
+     - If future requirements include bias/add/activation, they can be fused into the epilogue of this kernel to avoid extra memory traffic.
   
-     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.
+     Behavior:
+     - Accepts multiple calling conventions, including tuple-packed inputs.
+     - Validates shapes/dtypes/devices, allocates output (if needed), configures the Triton grid, and launches the kernel.
+     - All numerical computation is performed inside the Triton kernel using tl.load/tl.store/tl.dot; the wrapper does not perform math.
+  
+     Returns:
+     - Output tensor C with dtype matching the input tensors' dtype (fp16 in these tests).
      """
-     # 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).")
+     a, b, c = _extract_abc_from_args(*args)
   
-     # 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])}.")
+     if not isinstance(a, torch.Tensor) or not isinstance(b, torch.Tensor):
+         raise TypeError("A and B must be torch.Tensors.")
   
-     # Shapes and strides
-     M, K = a.shape
-     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):
+     if a.dim() != 2 or b.dim() != 2:
+         raise ValueError("A and B must be 2D matrices.")
+  
+     M, K_a = a.shape
+     K_b, N = b.shape
+     if K_a != K_b:
+         raise ValueError(f"Incompatible matmul shapes: {a.shape} @ {b.shape}")
+  
+     K = K_a
+  
+     if a.device.type != "cuda" or b.device.type != "cuda":
+         raise ValueError("A and B must be CUDA tensors.")
+  
+     if a.dtype != b.dtype:
+         raise ValueError("A and B must have the same dtype.")
+     if a.dtype not in (torch.float16, torch.bfloat16, torch.float32):
+         # Primary target is float16 as per tests; allow bf16/fp32 for completeness.
+         raise TypeError("Supported dtypes: float16, bfloat16, float32.")
+  
+     # Allocate output if not provided; use same dtype/device as A
+     if c is None:
          c = torch.empty((M, N), device=a.device, dtype=a.dtype)
+     else:
+         if not isinstance(c, torch.Tensor):
+             raise TypeError("Output buffer c must be a torch.Tensor if provided.")
+         if c.device != a.device:
+             raise ValueError("Output tensor device must match input tensors' device.")
+         if c.shape != (M, N):
+             raise ValueError(f"Output tensor shape must be {(M, N)} but got {tuple(c.shape)}.")
+         if c.dtype != a.dtype:
+             raise ValueError("Output tensor dtype must match input dtype.")
   
-     # Grid computation: 1D launch, flattened (pid_m, pid_n)
+     # Prepare strides for row-major tensors
+     stride_am, stride_ak = a.stride(0), a.stride(1)
+     stride_bk, stride_bn = b.stride(0), b.stride(1)
+     stride_cm, stride_cn = c.stride(0), c.stride(1)
+  
+     # Grid function uses autotuned BLOCK sizes
      def grid(meta):
-         return (triton.cdiv(M, meta["BLOCK_M"]) * triton.cdiv(N, meta["BLOCK_N"]),)
+         return (triton.cdiv(M, meta["BLOCK_M"]), triton.cdiv(N, meta["BLOCK_N"]))
   
      _matmul_kernel[grid](
          a, b, c,
          M, N, K,
-         a.stride(0), a.stride(1),
-         b.stride(0), b.stride(1),
-         c.stride(0), c.stride(1),
+         stride_am, stride_ak,
+         stride_bk, stride_bn,
+         stride_cm, stride_cn,
      )
-  
-     # Return the output tensor
      return c
   
  import inspect