code: def forward(
      ctx,
      a,
      scale,
      rel_logits,
      is_causal,
      spdims,
      block,
      lut,
      maxlut,
      is_dense,
  ):
      if scale is not None and isinstance(scale, torch.Tensor):
          assert scale.device.type == "cpu"
          scale = scale.item()
      M = a.shape[0]
      grid = [spdims[0], spdims[1] * block, M]
      rel_shape = (
          (1, 1, 1, 1) if rel_logits is None else rel_logits.shape
      )
      rel_strides = (
          (1, 1, 1, 1) if rel_logits is None else rel_logits.stride()
      )
      # enqueue kernel
      out = torch.empty_like(a)
      _blocksparse_softmax_fwd[grid](
          out,
          a,
          a.stride(0),
          lut,  #
          rel_logits,
          rel_shape[-1],
          rel_strides[0],
          rel_strides[1],  # relative attn#
          scale,  #
          is_causal,  #
          BLOCK_SIZE=block,  #
          ROW_SIZE=next_power_of_2(maxlut),  #
          IS_DENSE=is_dense,  #
          num_warps=num_warps(maxlut),  #
      )
      # save to context
      # ctx.mark_dirty(x)
      ctx.save_for_backward(out, lut)
      ctx.spdims = spdims
      ctx.block = block
      ctx.maxlut = maxlut
      ctx.scale = scale
      ctx.rel_shape = rel_shape
      ctx.rel_strides = rel_strides
      ctx.rel_dtype = a.dtype
      ctx.is_dense = is_dense
      ctx.is_causal = is_causal
      return out
  
file_path: triton/python/triton/ops/blocksparse/softmax.py
imports: import torch
  from ...None import jit
  from ...None import language
  from ...None import next_power_of_2
label_desc: Use image node_id 2 for calling the _softmax obj's underlying member method code with example usage: obj.forward(ctx, a, scale, rel_logits, is_causal, spdims, block, lut, maxlut, is_dense) and returns: out