code: def _forward(
      X,
      scale,
      LUT,
      RPE,
      KP_M,
      ATTN_M,
      sizemax,
      stride_zx,
      stride_zrpe,
      stride_hrpe,
      stride_srpe,
      stride_zkpm,
      stride_zattnm,
      **meta,
  ):
      TN = meta["TN"]
      BLOCK = meta["BLOCK"]
      pidhm = tl.program_id(0)
      pidz = tl.program_id(1)
      # create index ranges
      rxm = pidhm % BLOCK
      rbm = pidhm // BLOCK
      rxn = tl.arange(0, TN) % BLOCK
      rbn = tl.arange(0, TN) // BLOCK
      # extract information from LUT
      header = LUT + rbm * 2
      size = tl.load(header + 0)
      offset = tl.load(header + 1)
      check = rbn < size
      rbmn = tl.where(check, rbn, size - 1)
      # block id and column id
      blockid = tl.load(LUT + offset + rbmn * 4 + 0)
      columnid = tl.load(LUT + offset + rbmn * 4 + 1)
      rowid = tl.load(LUT + offset + rbmn * 4 + 2)
      headid = tl.load(LUT + offset + rbmn * 4 + 3)
      # pointers to X
      px = (
          X
          + pidz * stride_zx
          + blockid * BLOCK * BLOCK
          + rxm * BLOCK
          + rxn
      )
      x = tl.load(px, mask=check, other=-float("inf"))
      x = x.to(tl.float32)
      # apply scale
      if meta["APPLY_SCALE"]:
          x = x * scale
      # apply RPE
      if meta["APPLY_RPE"]:
          prpe = (
              RPE
              + pidz * stride_zrpe
              + headid * stride_hrpe
              + columnid * BLOCK
              + rowid * BLOCK * stride_srpe
              + rxm * stride_srpe
              + rxn
          )
          rpe = tl.load(prpe, mask=check, other=0)
          x = x + rpe
      # apply key-padding mask
      if meta["APPLY_KP_MASK"]:
          pkp_m = KP_M + pidz * stride_zkpm + columnid * BLOCK + rxn
          kp_m = tl.load(pkp_m, mask=check, other=-float("inf"))
          if meta["KP_MASK_MUL"]:
              kp_m = tl.where(kp_m == 0, -float("inf"), 0.0)
          x = x + kp_m
      # apply attention mask
      if meta["APPLY_ATTN_MASK"]:
          pattn_m = (
              ATTN_M
              + columnid * BLOCK
              + rowid * BLOCK * stride_zattnm
              + rxm * stride_zattnm
              + rxn
          )
          attn_m = tl.load(pattn_m, mask=check, other=-float("inf"))
          if meta["ATTN_MASK_MUL"]:
              attn_m = tl.where(attn_m == 0, -float("inf"), 0.0)
          x = x + attn_m
      # computation
      x = tl.softmax(x)
      tl.store(px, x, mask=check)
  
file_path: DeepSpeed/deepspeed/ops/sparse_attention/softmax.py
imports: import torch
  import triton
  import triton.language
label_desc: Use image node_id 3 for calling a global function with example usage: _forward(X, scale, LUT, RPE, KP_M, ATTN_M, sizemax, stride_zx, stride_zrpe, stride_hrpe, stride_srpe, stride_zkpm, stride_zattnm) without return types