class_docstr: Block-Sparse Softmax class; this class computes softmax on a block sparse matrix. It is also able to apply either/all of the following masks:
     - relative position embedding
     - key padding mask
     - attention mask
  
  For more details about sparsity config, please see `Generative Modeling with Sparse Transformers`: https://arxiv.org/abs/1904.10509
class_docstr_tok: ["Block-Sparse","Softmax","class",";","this","class","computes","softmax","on","a","block","sparse","matrix",".","It","is","also","able","to","apply","either\/all","of","the","following","masks",":","-","relative","position","embedding","-","key","padding","mask","-","attention","mask","For","more","details","about","sparsity","config",",","please","see","`","Generative","Modeling","with","Sparse","Transformers","`",":","https",":","\/\/arxiv.org\/abs\/1904.10509"]
docstr: Applies softmax on a Block-Sparse input tensor.
  
  For more details about sparsity config, please see `Generative Modeling with Sparse Transformers`: https://arxiv.org/abs/1904.10509
  
  Arguments:
       x: required: a block-sparse tensor that softmax is applied on it; computation will be in place and result will be returned in the same tensor
       scale: optional: a float value; x values will be multiplied by this value before normalization. Default value is 1.0.
       rpe: optional: a tensor same dimension as x that is used as relative position embedding
       key_padding_mask: optional: a mask tensor of size (BatchSize X SequenceLength)
       attn_mask: optional: a mask tensor of size (SequenceLength X SequenceLength); currently only 2D is supported
       key_padding_mask_mode: optional: a boolean determining if key_padding_mask needs to be added or multiplied
       attn_mask_mode: optional: a boolean determining if attn_mask needs to be added or multiplied
  
  Return:
       x: a block-sparse tensor contains normalized input x using softmax; and masks applied if given
docstr_tok: ["Applies","softmax","on","a","Block-Sparse","input","tensor",".","For","more","details","about","sparsity","config",",","please","see","`","Generative","Modeling","with","Sparse","Transformers","`",":","https",":","\/\/arxiv.org\/abs\/1904.10509","Arguments",":","x",":","required",":","a","block-sparse","tensor","that","softmax","is","applied","on","it",";","computation","will","be","in","place","and","result","will","be","returned","in","the","same","tensor","scale",":","optional",":","a","float","value",";","x","values","will","be","multiplied","by","this","value","before","normalization",".","Default","value","is","1.0.","rpe",":","optional",":","a","tensor","same","dimension","as","x","that","is","used","as","relative","position","embedding","key_padding_mask",":","optional",":","a","mask","tensor","of","size","(","BatchSize","X","SequenceLength",")","attn_mask",":","optional",":","a","mask","tensor","of","size","(","SequenceLength","X","SequenceLength",")",";","currently","only","2D","is","supported","key_padding_mask_mode",":","optional",":","a","boolean","determining","if","key_padding_mask","needs","to","be","added","or","multiplied","attn_mask_mode",":","optional",":","a","boolean","determining","if","attn_mask","needs","to","be","added","or","multiplied","Return",":","x",":","a","block-sparse","tensor","contains","normalized","input","x","using","softmax",";","and","masks","applied","if","given"]
code: def __call__(
      self,
      x,
      scale=1.0,
      rpe=None,
      key_padding_mask=None,
      attn_mask=None,
      key_padding_mask_mode="add",
      attn_mask_mode="add",
  ):
      """Applies softmax on a Block-Sparse input tensor.
  
      For more details about sparsity config, please see `Generative Modeling with Sparse Transformers`: https://arxiv.org/abs/1904.10509
  
      Arguments:
           x: required: a block-sparse tensor that softmax is applied on it; computation will be in place and result will be returned in the same tensor
           scale: optional: a float value; x values will be multiplied by this value before normalization. Default value is 1.0.
           rpe: optional: a tensor same dimension as x that is used as relative position embedding
           key_padding_mask: optional: a mask tensor of size (BatchSize X SequenceLength)
           attn_mask: optional: a mask tensor of size (SequenceLength X SequenceLength); currently only 2D is supported
           key_padding_mask_mode: optional: a boolean determining if key_padding_mask needs to be added or multiplied
           attn_mask_mode: optional: a boolean determining if attn_mask needs to be added or multiplied
  
      Return:
           x: a block-sparse tensor contains normalized input x using softmax; and masks applied if given
      """
  
      time_y = [None]
      if rpe is not None and rpe.dtype != x.dtype:
          raise ValueError(
              "relative position embedding must be %s" % x.dtype
          )
      if attn_mask is not None and attn_mask.dtype != x.dtype:
          raise ValueError("Attention mask must be %s" % x.dtype)
      if (
          key_padding_mask is not None
          and key_padding_mask.dtype != x.dtype
      ):
          raise ValueError("Key padding mask must be %s" % x.dtype)
      lut, maxlut = self.make_lut(x.device)
      x = Softmax.sparse_softmax(
          x,
          scale,
          rpe,
          key_padding_mask,
          attn_mask,
          key_padding_mask_mode,
          attn_mask_mode,
          self.spdims,
          self.block,
          lut,
          self.num_blocks,
          maxlut,
          self.bench,
          time_y,
      )
      self.time_y = time_y[0]
      return x
  
file_path: H2O/h2o_flexgen/benchmark/third_party/DeepSpeed/deepspeed/ops/sparse_attention/softmax.py
imports: import torch
  import triton
  import triton.language
label_desc: Use image node_id 4 for calling the Softmax obj's underlying member method code with example usage: obj.__call__(x, scale, rpe, key_padding_mask, attn_mask, key_padding_mask_mode, attn_mask_mode) and returns: x