submission 780703

  import torch
  import triton
  import triton.language as tl
+  
  from task import input_t, output_t
- from utils import make_match_reference, DeterministicContext
   
   
- def generate_input(m: int, n: int, k: int, seed: int) -> input_t:
-     gen = torch.Generator(device='cuda')
-     gen.manual_seed(seed)
-     a = torch.empty(m, k, device='cuda', dtype=torch.float16)
-     a.uniform_(0, 1, generator=gen)
-     b = torch.empty(k, n, device='cuda', dtype=torch.float16)
-     b.uniform_(0, 1, generator=gen)
-     c = torch.empty(m, n, device='cuda', dtype=torch.float16)
-     return a, b, c
-  
-  
- @triton.autotune(
-     configs=[
-         # Small problem sizes: finer tile granularity for better SM occupancy
-         triton.Config({'BLOCK_M': 64,  'BLOCK_N': 64,  'BLOCK_K': 64,  'GROUP_SIZE_M': 8}, num_stages=3, num_warps=4),
-         triton.Config({'BLOCK_M': 64,  'BLOCK_N': 128, 'BLOCK_K': 64,  'GROUP_SIZE_M': 8}, num_stages=3, num_warps=4),
-         triton.Config({'BLOCK_M': 128, 'BLOCK_N': 64,  'BLOCK_K': 64,  'GROUP_SIZE_M': 8}, num_stages=3, num_warps=4),
-         # Medium problem sizes: balanced tiles with moderate K
-         triton.Config({'BLOCK_M': 128, 'BLOCK_N': 128, 'BLOCK_K': 64,  'GROUP_SIZE_M': 8}, num_stages=3, num_warps=8),
-         triton.Config({'BLOCK_M': 128, 'BLOCK_N': 128, 'BLOCK_K': 128, 'GROUP_SIZE_M': 8}, num_stages=2, num_warps=8),
-         triton.Config({'BLOCK_M': 256, 'BLOCK_N': 128, 'BLOCK_K': 64,  'GROUP_SIZE_M': 8}, num_stages=3, num_warps=8),
-         triton.Config({'BLOCK_M': 128, 'BLOCK_N': 256, 'BLOCK_K': 64,  'GROUP_SIZE_M': 8}, num_stages=3, num_warps=8),
-         # Large problem sizes: wide tiles to maximize compute per block
-         triton.Config({'BLOCK_M': 256, 'BLOCK_N': 256, 'BLOCK_K': 64,  'GROUP_SIZE_M': 8}, num_stages=3, num_warps=8),
-         triton.Config({'BLOCK_M': 128, 'BLOCK_N': 256, 'BLOCK_K': 128, 'GROUP_SIZE_M': 8}, num_stages=2, num_warps=8),
-         # Deep K dimension: large K blocks to reduce loop iterations
-         triton.Config({'BLOCK_M': 128, 'BLOCK_N': 128, 'BLOCK_K': 256, 'GROUP_SIZE_M': 8}, num_stages=1, num_warps=8),
-         triton.Config({'BLOCK_M': 256, 'BLOCK_N': 128, 'BLOCK_K': 128, 'GROUP_SIZE_M': 8}, num_stages=2, num_warps=8),
-     ],
-     key=['M', 'N', 'K'],
- )
  @triton.jit
- def b200_matmul_kernel(
-     a_ptr, b_ptr, out_ptr,
-     M, N, K,
-     stride_am, stride_ak,
-     stride_bk, stride_bn,
-     stride_om, stride_on,
-     BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr,
-     GROUP_SIZE_M: tl.constexpr,
+ def _matmul_kernel(
+     a_ptr,
+     b_ptr,
+     c_ptr,
+     M,
+     N,
+     K,
+     stride_am,
+     stride_ak,
+     stride_bk,
+     stride_bn,
+     stride_cm,
+     stride_cn,
+     BLOCK_M: tl.constexpr,
+     BLOCK_N: tl.constexpr,
+     BLOCK_K: tl.constexpr,
  ):
-     """
-     High-performance FP16 matmul kernel tuned for NVIDIA B200 (Blackwell).
-  
-     Uses group-ordered launch for L2 cache optimization, tensor-core-backed
-     tl.dot with FP32 accumulation, and configurable tile sizes selected via
-     autotuning across the benchmark shape spectrum.
-     """
      pid = tl.program_id(0)
-  
-     # --- Grouped launch ordering ---
-     # Programs that share A-tile rows are grouped together to improve L2 hit rate.
-     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 = tl.minimum(num_pid_m - first_pid_m, GROUP_SIZE_M)
-     pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m)
-     pid_n = (pid % num_pid_in_group) // group_size_m
+     pid_m = pid // num_pid_n
+     pid_n = pid % num_pid_n
   
-     # --- Tile coordinate ranges ---
      offs_am = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
      offs_bn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
-     offs_k  = tl.arange(0, BLOCK_K)
+     offs_k = tl.arange(0, BLOCK_K)
   
-     # --- Pointer base addresses ---
      a_ptrs = a_ptr + offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak
      b_ptrs = b_ptr + offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn
   
-     accumulator = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+     acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
   
-     # --- Main K-loop ---
-     num_k_blocks = tl.cdiv(K, BLOCK_K)
-     for k in range(0, num_k_blocks):
-         k_rem = K - k * BLOCK_K
-         a = tl.load(a_ptrs, mask=offs_k[None, :] < k_rem, other=0.0)
-         b = tl.load(b_ptrs, mask=offs_k[:, None] < k_rem, other=0.0)
-         accumulator = tl.dot(a, b, accumulator, allow_tf32=False)
+     for k in range(0, K, BLOCK_K):
+         k_mask = (k + offs_k) < K
+         a = tl.load(a_ptrs, mask=k_mask[None, :], other=0.0)
+         b = tl.load(b_ptrs, mask=k_mask[:, None], other=0.0)
+         acc += tl.dot(a, b)
          a_ptrs += BLOCK_K * stride_ak
          b_ptrs += BLOCK_K * stride_bk
   
-     # --- Store result ---
-     out = accumulator.to(tl.float16)
-     offs_om = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
-     offs_on = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
-     out_ptrs = out_ptr + offs_om[:, None] * stride_om + offs_on[None, :] * stride_on
-     out_mask = (offs_om[:, None] < M) & (offs_on[None, :] < N)
-     tl.store(out_ptrs, out, mask=out_mask)
+     c = acc.to(tl.float16)
+     offs_cm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
+     offs_cn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
+     c_ptrs = c_ptr + offs_cm[:, None] * stride_cm + offs_cn[None, :] * stride_cn
+     c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N)
+     tl.store(c_ptrs, c, mask=c_mask)
   
   
  def custom_kernel(data: input_t) -> output_t:
-     a, b, _c = data
+     a, b, c = data
      M, K = a.shape
      K2, N = b.shape
-     assert K == K2, f"Inner dimension mismatch: {K} != {K2}"
   
-     output = torch.empty(M, N, device='cuda', dtype=torch.float16)
+     BLOCK_M = 128
+     BLOCK_N = 256
+     BLOCK_K = 64
   
-     grid = lambda meta: (
-         triton.cdiv(M, meta['BLOCK_M']) * triton.cdiv(N, meta['BLOCK_N']),
-     )
+     grid = (triton.cdiv(M, BLOCK_M) * triton.cdiv(N, BLOCK_N),)
   
-     b200_matmul_kernel[grid](
-         a, b, output,
-         M, N, K,
-         a.stride(0), a.stride(1),
-         b.stride(0), b.stride(1),
-         output.stride(0), output.stride(1),
+     _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),
+         BLOCK_M=BLOCK_M,
+         BLOCK_N=BLOCK_N,
+         BLOCK_K=BLOCK_K,
+         num_warps=8,
+         num_stages=3,
      )
-     return output
-  
-  
- check_implementation = make_match_reference(custom_kernel)
+     return c