code: def kernel_bwd(
      C,  # Pointers to matrices
      ACT_INPUT,
      A,
      B,
      # Matrix dimensions
      M,
      N,
      K,
      CACHE_KEY_M,
      CACHE_KEY_N,
      CACHE_KEY_K,
      # The stride variables represent how much to increase the ptr by when moving by 1
      # element in a particular dimension. E.g. stride_am is how much to increase a_ptr
      # by to get the element one row down (A has M rows)
      stride_cm,
      # stride_cn,  # Assume that stride_cn == 1
      stride_am,
      stride_ak,
      stride_bk,
      stride_bn,
      # Meta-parameters
      BLOCK_M: tl.constexpr,
      GROUP_M: tl.constexpr,
      BLOCK_N: tl.constexpr,
      BLOCK_K: tl.constexpr,
      # split k not used, not performant with activation, kept because early_config_prune is expecting it
      SPLIT_K: tl.constexpr,
      EVEN_K: tl.constexpr,
      ACTIVATION: tl.constexpr,
  ):
      """
      Kernel for computing Out = activation(A x W + C)
      - Input has shape (M, K)
      - Weight has shape (K, N)
      - Output has shape (M, N)
      - ActInputs (optional) has shape (M, N)
      'ActInputs' optionally saves the A x W + C intermediate for backward computations
      This kernel will consolidate over K
      """
  
      pid = tl.program_id(axis=0)
  
      grid_m = (M + BLOCK_M - 1) // BLOCK_M
      grid_n = (N + BLOCK_N - 1) // BLOCK_N
      # re-order program ID for better L2 performance
      width = GROUP_M * grid_n
      group_id = pid // width
      group_size = min(grid_m - group_id * GROUP_M, GROUP_M)
      pid_m = group_id * GROUP_M + (pid % group_size)
      pid_n = (pid % width) // (group_size)
  
      # now compute the block that each program will go through
      # rm (resp. rn) denotes a range of indices
      # for rows (resp. col) of C
      rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
      rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
      # trick to avoid masking on M and N axis
      ram = tl.max_contiguous(tl.multiple_of(rm % M, BLOCK_M), BLOCK_M)
      rbn = tl.max_contiguous(tl.multiple_of(rn % N, BLOCK_N), BLOCK_N)
      rk = tl.arange(0, BLOCK_K)
  
      A = A + (ram[:, None] * stride_am + rk[None, :] * stride_ak)
      B = B + (rk[:, None] * stride_bk + rbn[None, :] * stride_bn)
  
      acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
  
      for k in range(K, 0, -BLOCK_K):
          if EVEN_K:
              a = tl.load(A)
              b = tl.load(B)
          else:
              a = tl.load(A, mask=rk[None, :] < k, other=0.0)
              b = tl.load(B, mask=rk[:, None] < k, other=0.0)
          acc += tl.dot(a, b)
  
          A += BLOCK_K * stride_ak
          B += BLOCK_K * stride_bk
  
      # optional: fused activation (while the data is in shared memory)
      if ACTIVATION != "id":
          act_in_ptrs = (
              ACT_INPUT + ram[:, None] * stride_cm + rbn[None, :]
          )
          act_input = tl.load(act_in_ptrs).to(acc.dtype)
      if ACTIVATION == "gelu":
          acc *= gelu_grad(act_input)
      elif ACTIVATION == "gelu_approx":
          acc *= gelu_approx_grad(act_input)
      elif ACTIVATION == "squared_relu":
          acc *= squared_relu_grad(act_input)
  
      # rematerialize rm and rn to save registers
      rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
      rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
  
      # write back result
      C = C + rm[:, None] * stride_cm + rn[None, :]
      mask = (rm < M)[:, None] & (rn < N)[None, :]
      tl.store(C, acc, mask=mask)
  
file_path: flash-attention/flash_attn/ops/triton/linear.py
imports: from typing import Optional
  import torch
  import triton
  import triton.language
  from triton.ops.matmul_perf_model import early_config_prune, estimate_matmul_time
  from flash_attn.ops.triton.k_activations import gelu, gelu_approx, gelu_approx_grad, gelu_grad, squared_relu, squared_relu_grad
label_desc: Use image node_id 5 for calling a global function with example usage: kernel_bwd(C, ACT_INPUT, A, B, M, N, K, CACHE_KEY_M, CACHE_KEY_N, CACHE_KEY_K, stride_cm, stride_am, stride_ak, stride_bk, stride_bn, BLOCK_M, GROUP_M, BLOCK_N, BLOCK_K, SPLIT_K, EVEN_K, ACTIVATION) without return types