code: def _attn_fwd_inner(
      acc,
      l_i,
      m_i,
      q,  #
      K_block_ptr,
      V_block_ptr,  #
      start_m,
      qk_scale,  #
      BLOCK_M: tl.constexpr,
      BLOCK_DMODEL: tl.constexpr,
      BLOCK_N: tl.constexpr,  #
      STAGE: tl.constexpr,
      offs_m: tl.constexpr,
      offs_n: tl.constexpr,  #
      N_CTX: tl.constexpr,
  ):
      # range of values handled by this stage
      if STAGE == 1:
          lo, hi = 0, start_m * BLOCK_M
      elif STAGE == 2:
          lo, hi = start_m * BLOCK_M, (start_m + 1) * BLOCK_M
          lo = tl.multiple_of(lo, BLOCK_M)
      # causal = False
      else:
          lo, hi = 0, N_CTX
      K_block_ptr = tl.advance(K_block_ptr, (0, lo))
      V_block_ptr = tl.advance(V_block_ptr, (lo, 0))
      # loop over k, v and update accumulator
      for start_n in range(lo, hi, BLOCK_N):
          start_n = tl.multiple_of(start_n, BLOCK_N)
          # -- compute qk ----
          k = tl.load(K_block_ptr)
          qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
          qk += tl.dot(q, k)
          if STAGE == 2:
              mask = offs_m[:, None] >= (start_n + offs_n[None, :])
              qk = qk * qk_scale + tl.where(mask, 0, -1.0e6)
              m_ij = tl.maximum(m_i, tl.max(qk, 1))
              qk -= m_ij[:, None]
          else:
              m_ij = tl.maximum(m_i, tl.max(qk, 1) * qk_scale)
              qk = qk * qk_scale - m_ij[:, None]
          p = tl.math.exp2(qk)
          l_ij = tl.sum(p, 1)
          # -- update m_i and l_i
          alpha = tl.math.exp2(m_i - m_ij)
          l_i = l_i * alpha + l_ij
          # -- update output accumulator --
          acc = acc * alpha[:, None]
          # update acc
          v = tl.load(V_block_ptr)
          acc += tl.dot(p.to(tl.float16), v)
          # update m_i and l_i
          m_i = m_ij
          V_block_ptr = tl.advance(V_block_ptr, (BLOCK_N, 0))
          K_block_ptr = tl.advance(K_block_ptr, (0, BLOCK_N))
      return acc, l_i, m_i
  
file_path: triton/python/tutorials/06-fused-attention.py
imports: import pytest
  import torch
  import triton
  import triton.language
label_desc: Use image node_id 1 for calling a global function with example usage: _attn_fwd_inner(acc, l_i, m_i, q, K_block_ptr, V_block_ptr, start_m, qk_scale, BLOCK_M, BLOCK_DMODEL, BLOCK_N, STAGE, offs_m, offs_n, N_CTX) and returns: acc,  l_i,  m_i