code: def forward(
      self,
      query_key_value: torch.Tensor,
      attn_mask: torch.Tensor,
      heads: int,
      num_kv: int,
      norm_factor: float,
      no_masking: bool,
      layer_id: int,
      num_layers: int,
      alibi: torch.Tensor,
  ):
      if alibi is not None:
          batch_heads = query_key_value.shape[0] * heads
          offset = (
              dist.get_rank() * batch_heads
              if dist.is_initialized()
              else 0
          )
          alibi = alibi[offset : batch_heads + offset, :, :]
      else:
          alibi = torch.empty(1)
  
      output = self.softmax_context_func(
          query_key_value,
          attn_mask,
          self.config.rotary_dim,
          self.config.rotate_half,
          self.config.rotate_every_two,
          heads,
          num_kv,
          norm_factor,
          self.config.triangular_masking,
          self.config.local_attention,
          self.config.window_size,
          no_masking,
          layer_id,
          num_layers,
          alibi,
          self.config.rope_theta,
      )
  
      return output
  
file_path: deepspeed/deepspeed/ops/transformer/inference/op_binding/softmax_context.py
imports: import torch
  from deepspeed import comm
  from ..config import DeepSpeedInferenceConfig
  from .base import BaseOp
label_desc: Use image node_id 3 for calling the SoftmaxContextOp obj's underlying member method code with example usage: obj.forward(query_key_value, attn_mask, heads, num_kv, norm_factor, no_masking, layer_id, num_layers, alibi) and returns: output