- op : masked_multihead_attention_
  args : (Tensor x, Tensor cache_kv, Tensor bias, Tensor src_mask, Tensor cum_offsets, Tensor sequence_lengths, Tensor rotary_tensor, Tensor beam_cache_offset, Tensor qkv_out_scale, Tensor out_shift, Tensor out_smooth, int seq_len, int rotary_emb_dims, bool use_neox_rotary_style=false, str compute_dtype = "default", float out_scale=-1, int quant_round_type=1, float quant_max_bound=127.0, float quant_min_bound=-127.0)
  output : Tensor(out), Tensor(cache_kv_out), Tensor(beam_cache_offset_out)
  infer_meta :
    func : MaskedMultiheadAttentionInferMeta
  kernel :
    func : masked_multihead_attention
    data_type : x
  optional : bias, src_mask, cum_offsets, sequence_lengths, rotary_tensor, beam_cache_offset, qkv_out_scale, out_shift, out_smooth
  inplace : (cache_kv -> cache_kv_out), (beam_cache_offset -> beam_cache_offset_out)