code: def _cumulative_and_max_seq_len_nnz(
      qkv: torch.Tensor,
  ) -> Tuple[torch.Tensor, int, int]:
      # This function is used to calculate two pieces of metadata that are needed
      # for use with flash-attention and efficient_attention kernels. They are the
      # cumulative sequence_length over a batch of sequences and the maximum
      # sequence length.
  
      # It returns a tuple of cumulative sequence lengths and the maximum sequence
      # length, and the last element in the cumulative_sequence_lengths
      if not isinstance(qkv, NestedTensor):
          raise ValueError(
              "QKV must be nested for flash cumulative_seq_len calculation."
          )
  
      if qkv.lengths() is None:
          # TODO: Explore performance impact of copying
          cumulative_seqlen = qkv.offsets().to(
              dtype=torch.int32, device=qkv.device
          )
          max_seqlen = qkv._max_seqlen
          n_elem = qkv.values().shape[0]
      else:
          # TODO: Explore performance impact of copying
          cumulative_seqlen = (
              qkv.lengths()
              .cumsum(0)
              .to(dtype=torch.int32, device=qkv.device)
          )
          batch_size = qkv.size(0)
          max_seqlen = qkv._max_seqlen
          # TODO: Explore performance impact when compiling
          n_elem = int(cumulative_seqlen[-1].item())
      return cumulative_seqlen, max_seqlen, n_elem
  
imports: import logging
  import math
  from typing import Optional, Tuple
  import torch
  import torch.nn
  import torch.nn.functional
  from torch.backends.cuda import can_use_efficient_attention, can_use_flash_attention, flash_sdp_enabled, math_sdp_enabled, mem_efficient_sdp_enabled, SDPAParams, SDPBackend
  from .nested_tensor import buffer_from_jagged, NestedTensor, ViewNestedFromBuffer
label_desc: Use image node_id 11 for calling a global function with example usage: _cumulative_and_max_seq_len_nnz(qkv) and returns: cumulative_seqlen,  max_seqlen,  n_elem