submission 512161

+ # kernel.py
+ # Complete Triton implementation of a blocked matrix multiplication C = A @ B for fp16 tensors.
+ # The test harness imports and calls `kernel_function`, which launches the Triton kernel.
+ #
+ # Fusion note:
+ # - The test requires only A @ B. There is no bias or activation to fuse, so we implement a single-pass
+ #   matmul kernel with fp32 accumulation and fp16 store. If future requirements add bias, activation,
+ #   or epilogues, they can be fused into the same kernel to minimize memory traffic and launch overhead.
+  
  import torch
  import triton
  import triton.language as tl
  ⋯ 1 unchanged lines
   
  @triton.autotune(
      configs=[
-         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),
+         triton.Config({'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 64, 'BLOCK_SIZE_K': 32}, num_warps=4, num_stages=3),
+         triton.Config({'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 32}, num_warps=8, num_stages=4),
+         triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 64, 'BLOCK_SIZE_K': 32}, num_warps=8, num_stages=4),
+         triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 64}, num_warps=8, num_stages=4),
+         triton.Config({'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 256, 'BLOCK_SIZE_K': 32}, num_warps=8, num_stages=4),
+         triton.Config({'BLOCK_SIZE_M': 32, 'BLOCK_SIZE_N': 256, 'BLOCK_SIZE_K': 32}, num_warps=8, num_stages=4),
      ],
-     key=["M", "N", "K"],
+     key=['M', 'N', 'K'],
  )
  @triton.jit
  def _matmul_kernel(
  ⋯ 2 unchanged lines
      stride_am, stride_ak,
      stride_bk, stride_bn,
      stride_cm, stride_cn,
-     BLOCK_M: tl.constexpr,
-     BLOCK_N: tl.constexpr,
-     BLOCK_K: tl.constexpr,
+     BLOCK_SIZE_M: tl.constexpr,
+     BLOCK_SIZE_N: tl.constexpr,
+     BLOCK_SIZE_K: tl.constexpr,
  ):
-     """
-     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.
-     """
+     # Program IDs for the 2D launch grid
      pid_m = tl.program_id(axis=0)
      pid_n = tl.program_id(axis=1)
   
-     # 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_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)
+     # Compute the ranges this program instance will cover
+     start_m = pid_m * BLOCK_SIZE_M
+     start_n = pid_n * BLOCK_SIZE_N
   
-     # Accumulator in FP32
-     acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+     offs_m = start_m + tl.arange(0, BLOCK_SIZE_M)
+     offs_n = start_n + tl.arange(0, BLOCK_SIZE_N)
+     offs_k = tl.arange(0, BLOCK_SIZE_K)
   
-     # Iterate over K dimension
-     k_tiles = tl.cdiv(K, BLOCK_K)
-     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
+     # Provide alignment/contiguity hints to the compiler
+     offs_m = tl.max_contiguous(tl.multiple_of(offs_m, BLOCK_SIZE_M), BLOCK_SIZE_M)
+     offs_n = tl.max_contiguous(tl.multiple_of(offs_n, BLOCK_SIZE_N), BLOCK_SIZE_N)
+  
+     # Initialize accumulator in fp32 for better numerical accuracy
+     acc = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
+  
+     # Number of K-tiles to iterate
+     k_tiles = tl.cdiv(K, BLOCK_SIZE_K)
+  
+     # Loop over K dimension
+     for kt in range(k_tiles):
+         k_offset = kt * BLOCK_SIZE_K
+         # Compute pointers for A and B tiles
+         a_ptrs = a_ptr + (offs_m[:, None] * stride_am + (k_offset + offs_k[None, :]) * stride_ak)
+         b_ptrs = b_ptr + ((k_offset + offs_k[:, None]) * stride_bk + offs_n[None, :] * stride_bn)
+  
+         # Masks to guard OOB accesses
+         a_mask = (offs_m[:, None] < M) & ((k_offset + offs_k[None, :]) < K)
+         b_mask = ((k_offset + offs_k[:, None]) < K) & (offs_n[None, :] < N)
+  
+         # Load tiles
          a = tl.load(a_ptrs, mask=a_mask, other=0.0)
          b = tl.load(b_ptrs, mask=b_mask, other=0.0)
-         # Accumulate
+  
+         # Tensor core-friendly dot; acc is fp32
          acc = tl.dot(a, b, acc)
   
-     # Write back to C
-     c = acc.to(c_ptr.dtype.element_ty)
+     # Compute C pointers and mask
      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, c, mask=c_mask)
   
+     # Store the result to fp16
+     c_out = acc.to(tl.float16)
+     tl.store(c_ptrs, c_out, mask=c_mask)
   
- def _extract_abc_from_args(*args):
-     """
-     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
   
-  
- def kernel_function(*args):
+ def kernel_function(*args, **kwargs):
      """
-     Wrapper for a fused matmul kernel C = A @ B.
+     Python wrapper to launch the Triton matmul kernel C = A @ B.
   
-     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.
+     Accepted call patterns:
+     - kernel_function((a, b, c))
+     - kernel_function(a, b, c)
   
-     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.
+     Where:
+     - a: [M, K] fp16 CUDA tensor
+     - b: [K, N] fp16 CUDA tensor
+     - c: [M, N] fp16 CUDA tensor (output buffer). If provided, the kernel writes to it.
+       If not provided or mismatched, a new output tensor is allocated and returned.
   
+     Notes:
+     - All math is executed inside the Triton kernel. The wrapper performs validation,
+       allocation, and launch setup only, per runtime constraints.
+     - Fusion: Only matmul is required by the test. No additional ops to fuse. If epilogues
+       like bias or activation are needed later, they can be fused into the same kernel.
+  
      Returns:
-     - Output tensor C with dtype matching the input tensors' dtype (fp16 in these tests).
+     - The output tensor [M, N] fp16 on the same device as A.
      """
-     a, b, c = _extract_abc_from_args(*args)
+     # Unpack either tuple or separate args
+     if len(args) == 1 and isinstance(args[0], (tuple, list)) and len(args[0]) == 3:
+         a, b, c = args[0]
+     elif len(args) >= 3:
+         a, b, c = args[0], args[1], args[2]
+     else:
+         raise TypeError("kernel_function expects (a, b, c) either as a tuple/list or as separate arguments.")
   
-     if not isinstance(a, torch.Tensor) or not isinstance(b, torch.Tensor):
-         raise TypeError("A and B must be torch.Tensors.")
+     # Validate devices and dtypes: only CUDA fp16 as per test
+     if not isinstance(a, torch.Tensor) or not isinstance(b, torch.Tensor) or not isinstance(c, torch.Tensor):
+         raise TypeError("Arguments a, b, c 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 CUDA tensors.")
+     if a.dtype != torch.float16 or b.dtype != torch.float16 or c.dtype != torch.float16:
+         raise TypeError("All tensors must be float16 (fp16).")
   
-     if a.dim() != 2 or b.dim() != 2:
-         raise ValueError("A and B must be 2D matrices.")
-  
+     # Shape checks
+     if a.ndim != 2 or b.ndim != 2 or c.ndim != 2:
+         raise ValueError("a, b, c 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}")
-  
+         raise ValueError(f"Incompatible shapes: a is {a.shape}, b is {b.shape}. K must match.")
      K = K_a
+     if c.shape != (M, N):
+         # Allocate a new output if provided c has mismatch
+         c = torch.empty((M, N), device=a.device, dtype=torch.float16)
   
-     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.")
-  
-     # Prepare strides for row-major tensors
+     # Compute strides to support non-contiguous inputs
      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
+     # Grid calculation uses tl.cdiv via Triton's autotune grid function
      def grid(meta):
-         return (triton.cdiv(M, meta["BLOCK_M"]), triton.cdiv(N, meta["BLOCK_N"]))
+         return (triton.cdiv(M, meta['BLOCK_SIZE_M']), triton.cdiv(N, meta['BLOCK_SIZE_N']))
   
+     # Launch the Triton kernel
      _matmul_kernel[grid](
          a, b, c,
          M, N, K,
  ⋯ 1 unchanged lines
          stride_bk, stride_bn,
          stride_cm, stride_cn,
      )
+  
+     # Return the output buffer
      return c
   
  import inspect