code: def softmax_kernel(
      output_ptr,
      input_ptr,
      input_row_stride,
      output_row_stride,
      n_cols,
      BLOCK_SIZE: tl.constexpr,
  ):
      # The rows of the softmax are independent, so we parallelize across those
      row_idx = tl.program_id(0)
      # The stride represents how much we need to increase the pointer to advance 1 row
      row_start_ptr = input_ptr + row_idx * input_row_stride
      # The block size is the next power of two greater than n_cols, so we can fit each
      # row in a single block
      col_offsets = tl.arange(0, BLOCK_SIZE)
      input_ptrs = row_start_ptr + col_offsets
      # Load the row into SRAM, using a mask since BLOCK_SIZE may be > than n_cols
      row = tl.load(
          input_ptrs, mask=col_offsets < n_cols, other=-float("inf")
      )
      # Subtract maximum for numerical stability
      row_minus_max = row - tl.max(row, axis=0)
      # Note that exponentiation in Triton is fast but approximate (i.e., think __expf in CUDA)
      numerator = tl.exp(row_minus_max)
      denominator = tl.sum(numerator, axis=0)
      softmax_output = numerator / denominator
      # Write back output to DRAM
      output_row_start_ptr = output_ptr + row_idx * output_row_stride
      output_ptrs = output_row_start_ptr + col_offsets
      tl.store(output_ptrs, softmax_output, mask=col_offsets < n_cols)
  
imports: import torch
  import triton
  import triton.language
label_desc: Use image node_id 2 for calling a global function with example usage: softmax_kernel(output_ptr, input_ptr, input_row_stride, output_row_stride, n_cols, BLOCK_SIZE) without return types