submission 577884

  # kernel.py
+ # -----------------------------------------------------------------------------
+ # Custom Triton-based matmul (float16) implementing C = A @ B
+ #
+ # - Wrapper entry point: kernel_function
+ # - Accepts either a single tuple (a, b, c) or three separate args a, b, c
+ # - All math is performed inside the Triton kernel using tl.dot
+ # - Wrapper only validates inputs, allocates output if needed, and launches
+ # - Dimensions may be any size; masks handle edges. Tests use multiples of 16.
+ #
+ # Fusion note:
+ #   The requested pipeline is a pure matmul. There are no additional stages
+ #   (bias, activation, normalization) provided in the interface, so there is
+ #   nothing to fuse. If a bias/activation tensor were available, we would fuse
+ #   the epilogue (accumulator -> bias add -> activation) into this single pass.
+ # -----------------------------------------------------------------------------
+  
  import torch
  import triton
  import triton.language as tl
   
- """
- Triton matmul kernel and Python wrapper.
   
- What is implemented and why:
- - A blocked GEMM C = A @ B for 2D tensors, using fp32 accumulation and fp16 output.
- - The kernel follows Triton best practices:
-   - tl.load/tl.store with proper masks for boundary safety.
-   - Coalesced memory access along the contiguous dimension of each operand.
-   - 2D grid over (M, N) tiles and K-loop with BLOCK_K.
-   - Autotune over several tile shapes and pipeline stages/warps.
- - Epilogue fusion: convert accumulator (fp32) to output dtype (fp16) in-kernel and store.
-   No additional ops (bias/activation) are fused because the input contract only provides (A, B, C).
-   If such parameters existed, we would fuse them here to minimize memory traffic.
-  
- Runtime restrictions are respected:
- - The Python wrapper only validates inputs, prepares output storage, computes the grid,
-   and launches the Triton kernel. It does not perform any math (no torch.matmul, etc.).
- - All computation happens inside the Triton kernel via tl.dot and other Triton ops.
-  
- The wrapper accepts both signatures used by the tests:
- - kernel_function((a, b, c))
- - kernel_function(a, b, c)
- It returns the output tensor. If a third tensor c is provided, the kernel writes into it.
- """
-  
-  
- # Autotune configurations: cover common GEMM tile sizes with power-of-2 blocks
- # This selection aims to handle the provided test shapes efficiently and robustly.
+ # Autotune configurations: choose tile sizes and execution parameters.
+ # All BLOCK_* are powers of two for better performance on tensor cores.
+ # Accumulation is done in fp32 for improved fp16 numerical accuracy.
  _matmul_configs = [
-     triton.Config({"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 32}, num_warps=4, num_stages=3),
-     triton.Config({"BLOCK_M": 64,  "BLOCK_N": 128, "BLOCK_K": 32}, num_warps=4, num_stages=3),
-     triton.Config({"BLOCK_M": 128, "BLOCK_N": 64,  "BLOCK_K": 32}, num_warps=4, num_stages=3),
-     triton.Config({"BLOCK_M": 64,  "BLOCK_N": 256, "BLOCK_K": 32}, num_warps=8, num_stages=4),
-     triton.Config({"BLOCK_M": 64,  "BLOCK_N": 64,  "BLOCK_K": 64}, num_warps=4, num_stages=3),
-     triton.Config({"BLOCK_M": 128, "BLOCK_N": 256, "BLOCK_K": 32}, num_warps=8, num_stages=4),
+     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=8,
+     ),
+     triton.Config(
+         {"BLOCK_M": 128, "BLOCK_N": 64, "BLOCK_K": 32},
+         num_stages=3,
+         num_warps=8,
+     ),
+     triton.Config(
+         {"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 32},
+         num_stages=4,
+         num_warps=8,
+     ),
+     triton.Config(
+         {"BLOCK_M": 64, "BLOCK_N": 256, "BLOCK_K": 64},
+         num_stages=4,
+         num_warps=8,
+     ),
+     triton.Config(
+         {"BLOCK_M": 128, "BLOCK_N": 256, "BLOCK_K": 64},
+         num_stages=4,
+         num_warps=8,
+     ),
  ]
   
   
  ⋯ 5 unchanged lines
      stride_am, stride_ak,
      stride_bk, stride_bn,
      stride_cm, stride_cn,
-     # compile-time constants
-     BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr,
+     BLOCK_M: tl.constexpr,
+     BLOCK_N: tl.constexpr,
+     BLOCK_K: tl.constexpr,
  ):
-     # Program IDs for 2D launch: each program computes one [BLOCK_M x BLOCK_N] tile of C
+     """
+     Compute C = A @ B where:
+       - A is [M, K]
+       - B is [K, N]
+       - C is [M, N]
+     All tensors are assumed to be float16. Accumulation is in float32.
+     """
+     # Program IDs for the 2D launch grid
      pid_m = tl.program_id(axis=0)
      pid_n = tl.program_id(axis=1)
   
-     # Compute tile start offsets
+     # Tile indices for this program
      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)
   
-     # Improve codegen by hinting alignment/contiguity on offsets
-     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)
-     offs_k = tl.max_contiguous(tl.multiple_of(offs_k, BLOCK_K), BLOCK_K)
-  
-     # Create pointer grids for the first K tile. We'll update them in the K loop.
-     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)
-  
-     # Accumulator in float32 for better precision
+     # Create accumulator in fp32
      acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
   
-     # Number of K tiles
-     k_tiles = tl.cdiv(K, BLOCK_K)
-     # Mask helpers
-     m_mask = offs_m[:, None] < M            # broadcast across N
-     n_mask = offs_n[None, :] < N            # broadcast across M
-     # Loop over K tiles
-     for ki in range(0, k_tiles):
-         k_off = ki * BLOCK_K
-         # Update K pointers
-         a_tile_ptrs = a_ptrs + k_off * stride_ak
-         b_tile_ptrs = b_ptrs + k_off * stride_bk
+     # Loop over K dimension by BLOCK_K
+     # tl.range is used for a dynamic loop bound; it is lowered by Triton.
+     for k0 in tl.range(0, K, BLOCK_K):
+         k_idx = k0 + offs_k
   
-         # Mask for valid K indices in this tile
-         k_mask_row = (k_off + offs_k[None, :]) < K  # shape [1, BLOCK_K], broadcast over M
-         k_mask_col = (k_off + offs_k[:, None]) < K  # shape [BLOCK_K, 1], broadcast over N
+         # Pointers for A and B blocks
+         a_ptrs = a_ptr + (offs_m[:, None] * stride_am + k_idx[None, :] * stride_ak)
+         b_ptrs = b_ptr + (k_idx[:, None] * stride_bk + offs_n[None, :] * stride_bn)
   
-         # Load A and B tiles with masking (out-of-bounds elements are zero)
-         a = tl.load(a_tile_ptrs, mask=m_mask & k_mask_row, other=0.0)
-         b = tl.load(b_tile_ptrs, mask=k_mask_col & n_mask, other=0.0)
+         # Boundary masks
+         a_mask = (offs_m[:, None] < M) & (k_idx[None, :] < K)
+         b_mask = (k_idx[:, None] < K) & (offs_n[None, :] < N)
   
-         # Accumulate partial products
+         # Load tiles from A and B in fp16
+         a = tl.load(a_ptrs, mask=a_mask, other=0.0)
+         b = tl.load(b_ptrs, mask=b_mask, other=0.0)
+  
+         # Accumulate using tensor cores when available
          acc = tl.dot(a, b, acc)
   
-     # Write back: convert accumulator to output dtype (fp16 expected) and store with mask
+     # Write back results to C (cast to fp16)
      c_ptrs = c_ptr + (offs_m[:, None] * stride_cm + offs_n[None, :] * stride_cn)
-     c_mask = m_mask & n_mask
-     c_out = acc.to(c_ptr.dtype.element_ty)
-     tl.store(c_ptrs, c_out, mask=c_mask)
+     c_mask = (offs_m[:, None] < M) & (offs_n[None, :] < N)
+     tl.store(c_ptrs, acc.to(tl.float16), mask=c_mask)
   
   
- def kernel_function(*args):
+ def _parse_args(*args):
      """
-     Launch wrapper for the Triton GEMM kernel.
-  
-     Supports two call signatures:
-       - kernel_function((a, b, c))
-       - kernel_function(a, b, c)
-  
-     Arguments:
-       a: [M, K] input matrix on CUDA
-       b: [K, N] input matrix on CUDA
-       c: [M, N] output buffer on CUDA (optional; if provided, kernel writes to it)
-  
-     Returns:
-       Tensor [M, N] on CUDA with dtype matching inputs (fp16 recommended).
-  
-     Notes:
-       - All math is done inside the Triton kernel.
-       - The wrapper validates inputs, allocates output if needed, sets the grid, and launches.
+     Helper to accept either:
+       - a single tuple/list (a, b, c)
+       - or separate args (a, b, c)
+     Returns (a, b, c).
      """
-     # Unpack arguments from either signature
-     if len(args) == 1 and isinstance(args[0], tuple):
+     if len(args) == 1 and isinstance(args[0], (tuple, list)):
          if len(args[0]) != 3:
-             raise TypeError("Expected a 3-tuple (a, b, c)")
+             raise TypeError("Expected a tuple/list of (a, b, c).")
          a, b, c = args[0]
      elif len(args) == 3:
          a, b, c = args
      else:
-         raise TypeError("kernel_function expects (a, b, c) or ((a, b, c),)")
+         raise TypeError("kernel_function expects (a, b, c) either as a tuple or as separate args.")
+     return a, b, c
   
-     # Basic validation and setup
-     if not isinstance(a, torch.Tensor) or not isinstance(b, torch.Tensor):
-         raise TypeError("Inputs a and b must be torch.Tensors")
-     if a.dim() != 2 or b.dim() != 2:
-         raise ValueError("Inputs a and b must be 2D matrices")
-     if a.shape[1] != b.shape[0]:
-         raise ValueError(f"Incompatible shapes: A{a.shape} and B{b.shape}")
-     if a.device.type != "cuda" or b.device.type != "cuda":
-         raise ValueError("Inputs must be on CUDA")
   
-     M, K = a.shape
-     Kb, N = b.shape
-     assert K == Kb
+ def kernel_function(*args):
+     """
+     Launch wrapper for the Triton matmul kernel.
   
-     # Output allocation: use provided c if compatible; otherwise allocate a new buffer
-     out_dtype = a.dtype  # tests use fp16; we keep dtype consistency
-     if c is None or not isinstance(c, torch.Tensor):
-         c = torch.empty((M, N), device=a.device, dtype=out_dtype)
-     else:
-         if c.shape != (M, N):
-             raise ValueError(f"Provided c has shape {c.shape}, expected {(M, N)}")
-         if c.device != a.device:
-             raise ValueError("Output tensor c must be on the same CUDA device as inputs")
-         if c.dtype != out_dtype:
-             # Allow dtype mismatch but warn; test is tolerant if values match after cast
-             pass
+     Inputs:
+       - a: [M, K], float16, CUDA
+       - b: [K, N], float16, CUDA
+       - c: [M, N], float16, CUDA (output buffer; will be written in-place)
   
-     # Strides (support both contiguous and non-contiguous tensors)
+     Behavior:
+       - Computes c = a @ b using a single Triton kernel launch.
+       - Returns the output tensor 'c'. If desired, the caller may ignore the
+         return value and rely on the in-place write to 'c'.
+  
+     Notes and constraints:
+       - Wrapper performs only validation, allocation (not needed since c is provided),
+         and kernel launch. No numerical computation is done here.
+       - Accumulation happens in fp32 inside the kernel for better precision,
+         then converted to fp16 on store.
+       - All boundary conditions are masked. Tests use multiples of 16, but the
+         kernel is robust for arbitrary sizes.
+  
+     Fusion:
+       - Only matmul is requested. There are no compatible additional ops
+         provided to fuse (e.g., bias or activation), so this is already a
+         single-pass fused implementation for the requested workload.
+     """
+     a, b, c = _parse_args(*args)
+  
+     # Basic validation (setup-only; no compute)
+     if not (isinstance(a, torch.Tensor) and isinstance(b, torch.Tensor) and isinstance(c, torch.Tensor)):
+         raise TypeError("a, b, c must be torch.Tensor")
+  
+     if a.device.type != "cuda" or b.device.type != "cuda" or c.device.type != "cuda":
+         raise ValueError("All inputs 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 dtype torch.float16")
+  
+     if a.ndim != 2 or b.ndim != 2 or c.ndim != 2:
+         raise ValueError("a, b, c must be 2D tensors")
+  
+     M, K_a = a.shape
+     K_b, N = b.shape
+     Mc, Nc = c.shape
+  
+     if K_a != K_b:
+         raise ValueError(f"Incompatible shapes: A is {a.shape}, B is {b.shape}")
+     if Mc != M or Nc != N:
+         raise ValueError(f"Output C has wrong shape: expected {(M, N)}, got {c.shape}")
+  
+     # Strides in elements (PyTorch gives strides in elements already)
      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)
   
-     # 2D launch grid: one program per [BLOCK_M x BLOCK_N] tile of C
+     # Grid dimensions depend on autotuned tile 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"]),
+         )
   
-     # Launch Triton kernel
+     # Launch kernel
      _matmul_kernel[grid](
          a, b, c,
-         M, N, K,
+         M, N, K_a,
          stride_am, stride_ak,
          stride_bk, stride_bn,
          stride_cm, stride_cn,
      )
   
-     # Return the output tensor (the kernel already wrote into c)
+     # Return the output tensor (also written in-place)
      return c
   
  import inspect