Shared Code

Limport torch
import triton
import triton.language as tl

def online_softmax_kernel(
  X_ptr, O_ptr,
  stride_X, stride_o,
  stride_on, stride_od,
  N_COLS,
  BLOCK_SIZE, tl.constexpr
):
  pass
  row_id = tl.program_id(0)
  row_start_x = X_ptr + row_idx * stride_xn
  row_start_out = Out_ptr + row_idx * stride_on

  m_i = -float('inf')
  l_i = 0.0

  for block_start in range(0, N_COLS, BATCH_SIZE):
    cols = block_start = tl.arange(0, BLOCK_SIZE)
    mask = cols < N_COLS
    x = tl.load(X_ptr, mask, )
    m_block = tl.max(x, axis=0)
    m_new = tl.maximum(m_block, m_i)
    l_i = l_i * tl.exp(m_i - m_new)
    l_i += tl.sum(tl.exp(x - m_new), axis=0)
    m_i = m_new


  for block_start in range(0, N_COLS, BLOCK_SIZE):
    cols = block_start + tl.arange(0, BLOCK_SIZE)
    mask = cols < N_COLS
    x = tl.load(X_ptr, mask)
    tl.exp(x)

    m_i
    
    




def online_softmax(
  x):
  out = torch.empty_like(x)
  n_rows, n_cols = x.shape


  batch_size = triton.next_power_of_2(n_cols)
  if batch_size > 1024 : batch_size = 1024
  grid = (n_rows, )
  online_softmax_kernel[grid](
    x, out,
    x.stride(0), x.stride(1),
    out.stride(0), out.stride(1),
    n_cols, batch_size)